jax-ml/jax · error · NotImplementedError

f16/bf16 not supported for async atomics

Error message

f16/bf16 not supported for async atomics

What it means

f16/bf16 atomics in PTX are synchronous red/atom instructions; there is no async (::complete_tx::bytes) f16/bf16 atomic form to pair with a cluster barrier, so passing cluster_barrier_ptr with an f16/bf16 atomic store is rejected.

Source

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

      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}"
        )
      ptx_type = f"{element_type}x2"
      noftz = "" if multimem else ".noftz"
    else:
      raise NotImplementedError(
          f"Unsupported element type for atomic stores: {element_type}"
      )
    [vec_len] = vreg.type.shape
    if element_bitwidth == 16:
      if vec_len % 2 != 0:
        raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit cluster_barrier_ptr for the f16/bf16 atomic store
2. If asynchrony is required, accumulate in a wider type (i32) and convert
3. Keep the barrier for the non-atomic stores and issue the atomic store separately without it

Example fix

# before
fa16.store_tiled_async(ref, atomic='add', cluster_barrier_ptr=bar)
# after
fa16.store_tiled_async(ref, atomic='add')

from jax._src.lib import ir
if isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)) and cluster_barrier_ptr is not None:
    cluster_barrier_ptr = None
fa.store_tiled_async(ref, atomic=atomic, cluster_barrier_ptr=cluster_barrier_ptr)

from jax._src.lib import ir

def half_needs_sync(fa, barrier) -> bool:
    return isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)) and barrier is not None

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

When it happens

Trigger: Calling store_tiled_async on an F16Type/BF16Type array with atomic in ('add','min','max') and a non-None cluster_barrier_ptr.

Common situations: Low-precision accumulation pipelines (common on Hopper tensor-core kernels) where the async barrier plumbing from TMA stores is reused for the atomic accumulation store.

Related errors

• f16/bf16 only supports add, min, max atomics, got {atomic}
• f16/bf16 SMEM/multimem atomics only support add, got {atomic
• f32 not supported for async atomics
• f32 only supports add atomics, 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/a0b1193d42db3aa4. Report an issue: GitHub.