jax-ml/jax · error · NotImplementedError

f32 not supported for async atomics

Error message

f32 not supported for async atomics

What it means

For f32 atomic stores with a cluster barrier, Mosaic would need an async red/atom PTX form with ::complete_tx semantics, and no such f32 instruction exists, so combining cluster_barrier_ptr with an f32 atomic raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3893

      space = ""
      ptr_constraint = "l"
    elif multimem:
      assert not is_smem
      red = "multimem.red"
      scope = "sys"
      space = ".global"
      ptr_constraint = "l"
    else:
      red = "red"
      scope = "cta" if is_smem else "gpu"
      space = ".shared::cta" if is_smem else ""
      ptr_constraint = "r" if is_smem else "l"
    element_type = self.mlir_dtype
    element_bitwidth = utils.bitwidth(element_type)
    noftz = ""
    if isinstance(element_type, ir.F32Type):
      if cluster_barrier_ptr is not None:
        raise NotImplementedError("f32 not supported for async atomics")
      if atomic != "add":
        raise NotImplementedError(f"f32 only supports add atomics, got {atomic}")
      ptx_type = "f32"
    elif isinstance(element_type, ir.IntegerType) and element_bitwidth == 32:
      if atomic in ("and", "or", "xor"):
        ptx_type = "b32"
      else:
        ptx_type = "s32" if self.is_signed else "u32"
    elif isinstance(element_type, (ir.F16Type, ir.BF16Type)):
      if cluster_barrier_ptr is not None:
        raise NotImplementedError("f16/bf16 not supported for async atomics")
      if atomic not in ("add", "min", "max"):
        raise NotImplementedError(
            f"f16/bf16 only supports add, min, max atomics, got {atomic}"
        )
      if (is_smem or multimem) and atomic != "add":
        raise NotImplementedError(
            f"f16/bf16 SMEM/multimem atomics only support add, got {atomic}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop cluster_barrier_ptr for the f32 atomic store (make it synchronous)
2. Use i32/f32x2-packing tricks: store f32 bits via an integer-typed view if the backend supports it
3. Accumulate into f16/bf16 or integer shared-memory and convert afterwards, or perform the f32 atomic with a regular atomic op outside store_tiled_async

Example fix

// before
fa32.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=bar)
// after
fa32.store_tiled_async(ref, atomic='add')  # no async barrier for f32

from jax._src.lib import ir
if isinstance(fa.mlir_dtype, ir.F32Type) and cluster_barrier_ptr is not None:
    cluster_barrier_ptr = None  # f32 atomics are sync-only
fa.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=cluster_barrier_ptr)

from jax._src.lib import ir

def f32_atomic_needs_sync(fa, barrier) -> bool:
    return isinstance(fa.mlir_dtype, ir.F32Type) and barrier is not None

try:
    fa.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=bar)
except NotImplementedError:
    fa.store_tiled_async(ref, atomic='add')

When it happens

Trigger: Calling store_tiled_async with atomic='add' on an f32 FragmentedArray while also passing cluster_barrier_ptr (the TMA-style async completion barrier).

Common situations: Wiring an async pipeline with mbarrier complete_tx tracking and then performing f32 atomic accumulation; works for other dtypes so developers assume f32 is fine.

Related errors

• f32 only supports add atomics, got {atomic}
• f16/bf16 not supported for async atomics
• f16/bf16 only supports add, min, max atomics, got {atomic}
• f16/bf16 SMEM/multimem atomics only support add, got {atomic
• Unsupported element type for atomic stores: {element_type}

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