{"record":{"id":"53929a890e7071d1","repo":"jax-ml/jax","slug":"f32-redux-only-supported-on-blackwell-gpus","errorCode":null,"errorMessage":"F32 redux only supported on Blackwell GPUs","messagePattern":"F32 redux only supported on Blackwell GPUs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":2035,"sourceCode":"def redux(x: ir.Value, mask: ir.Value, kind: ReductionKind):\n  i32 = ir.IntegerType.get_signless(32)\n  if isinstance(vec_ty := x.type, ir.VectorType):\n    if bitwidth(vec_ty.element_type) != 32:\n      raise ValueError(\"Only 32-bit types supported\")\n    [vec_len] = vec_ty.shape\n    result = llvm.mlir_undef(x.type)\n    for i in range(vec_len):\n      xi = llvm.extractelement(x, arith.constant(i32, i))\n      yi = redux(xi, mask, kind)\n      result = llvm.insertelement(result, yi, arith.constant(i32, i))\n    return result\n  if bitwidth(x.type) != 32:\n    raise ValueError(\"Only 32-bit scalar types supported\")\n  if isinstance(x.type, ir.IntegerType):\n    pass\n  elif isinstance(x.type, ir.F32Type):\n    if get_arch().major != 10:\n      raise ValueError(\"F32 redux only supported on Blackwell GPUs\")\n  else:\n    raise NotImplementedError(x.type)\n  assert mask.type == i32\n  extra_kwargs: dict[str, Any] = {}\n  if kind == ReductionKind.FMAX or kind == ReductionKind.FMIN:\n    extra_kwargs = dict(nan=True)\n  return nvvm.redux_sync(x, kind, mask, **extra_kwargs)\n\n\ndef prmt(high: ir.Value, low: ir.Value, permutation: ir.Value):\n  i32 = ir.IntegerType.get_signless(32)\n  if (result_type := high.type) != low.type:\n    raise ValueError(f\"Types must match, got {high.type} and {low.type}\")\n  if high.type != i32:\n    high = bitcast(high, i32)\n  if low.type != i32:\n    low = bitcast(low, i32)\n  if permutation.type != i32:","sourceCodeStart":2017,"sourceCodeEnd":2053,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L2017-L2053","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Target a Blackwell GPU (set CUDA_VISIBLE_DEVICES or run on a B200/sm100 machine)","Branch on get_arch() and use warp_reduce/shuffle-based reduction for f32 on pre-Blackwell GPUs","Keep integer reductions (always supported) by bitcasting f32 to i32 only when ordering semantics allow — otherwise avoid"],"exampleFix":"# before\nres = redux(x_f32, mask, ReductionKind.FMIN)\n# after\nif get_arch().major == 10:\n  res = redux(x_f32, mask, ReductionKind.FMIN)\nelse:\n  res = warp_reduce(x_f32, mask, ReductionKind.FMIN)  # shuffle fallback","handlingStrategy":"fallback","validationCode":"from jax.experimental.mosaic.gpu.utils import get_arch\nif bitwidth(x.type) == 32 and isinstance(x.type, ir.F32Type):\n    assert get_arch().major == 10, 'f32 redux requires Blackwell (sm100)'","typeGuard":null,"tryCatchPattern":"try:\n    res = redux(x, mask, kind)\nexcept ValueError:\n    res = warp_reduce_fallback(x, mask, kind)  # shuffle-based, works on all archs","preventionTips":["Branch on get_arch() in kernels that must be portable","Pin run environments to the GPU arch the kernel was written for"],"tags":["mosaic-gpu","redux","gpu-architecture","blackwell"],"backgroundTag":"unsupported-gpu-architecture","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}