jax-ml/jax · error · ValueError

F32 redux only supported on Blackwell GPUs

Error message

F32 redux only supported on Blackwell GPUs

What it means

redux only accepts f32 floating-point values on SM100 (Blackwell), where PTX redux.sync with f32 operands was introduced. get_arch().major != 10 means an older architecture (Hopper SM90, Ada SM89, Ampere SM80...), so the f32 path raises.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2035

def redux(x: ir.Value, mask: ir.Value, kind: ReductionKind):
  i32 = ir.IntegerType.get_signless(32)
  if isinstance(vec_ty := x.type, ir.VectorType):
    if bitwidth(vec_ty.element_type) != 32:
      raise ValueError("Only 32-bit types supported")
    [vec_len] = vec_ty.shape
    result = llvm.mlir_undef(x.type)
    for i in range(vec_len):
      xi = llvm.extractelement(x, arith.constant(i32, i))
      yi = redux(xi, mask, kind)
      result = llvm.insertelement(result, yi, arith.constant(i32, i))
    return result
  if bitwidth(x.type) != 32:
    raise ValueError("Only 32-bit scalar types supported")
  if isinstance(x.type, ir.IntegerType):
    pass
  elif isinstance(x.type, ir.F32Type):
    if get_arch().major != 10:
      raise ValueError("F32 redux only supported on Blackwell GPUs")
  else:
    raise NotImplementedError(x.type)
  assert mask.type == i32
  extra_kwargs: dict[str, Any] = {}
  if kind == ReductionKind.FMAX or kind == ReductionKind.FMIN:
    extra_kwargs = dict(nan=True)
  return nvvm.redux_sync(x, kind, mask, **extra_kwargs)


def prmt(high: ir.Value, low: ir.Value, permutation: ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if (result_type := high.type) != low.type:
    raise ValueError(f"Types must match, got {high.type} and {low.type}")
  if high.type != i32:
    high = bitcast(high, i32)
  if low.type != i32:
    low = bitcast(low, i32)
  if permutation.type != i32:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Target a Blackwell GPU (set CUDA_VISIBLE_DEVICES or run on a B200/sm100 machine)
2. Branch on get_arch() and use warp_reduce/shuffle-based reduction for f32 on pre-Blackwell GPUs
3. Keep integer reductions (always supported) by bitcasting f32 to i32 only when ordering semantics allow — otherwise avoid

Example fix

# before
res = redux(x_f32, mask, ReductionKind.FMIN)
# after
if get_arch().major == 10:
  res = redux(x_f32, mask, ReductionKind.FMIN)
else:
  res = warp_reduce(x_f32, mask, ReductionKind.FMIN)  # shuffle fallback

from jax.experimental.mosaic.gpu.utils import get_arch
if bitwidth(x.type) == 32 and isinstance(x.type, ir.F32Type):
    assert get_arch().major == 10, 'f32 redux requires Blackwell (sm100)'

try:
    res = redux(x, mask, kind)
except ValueError:
    res = warp_reduce_fallback(x, mask, kind)  # shuffle-based, works on all archs

When it happens

Trigger: Calling redux on an f32 value while targeting any GPU architecture whose compute capability major version is not 10 (e.g. H100 SM90, A100 SM80).

Common situations: Developing a Mosaic kernel on Blackwell and running it on a Hopper or older cluster; JAX selecting a different XLA backend/GPU than expected; CI machines with older GPUs.

Related errors

• Stores to TMEM are asynchronous operations and cannot be per
• copy_gmem_to_smem with a barrier is only supported Hopper an
• Only the TMA implementation supports collective copies
• Only the TMA implementation supports leader_tracked copies
• The cp.async implementation only supports oob_mode=OOBFillMo

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/53929a890e7071d1. Report an issue: GitHub.