jax-ml/jax · error · NotImplementedError

f16/bf16 SMEM/multimem atomics only support add, got {atomic

Error message

f16/bf16 SMEM/multimem atomics only support add, got {atomic}

What it means

For f16/bf16, only the packed .add.f16x2 red form is available for shared-memory (DSMEM) and multimem destinations; min/max exist only for global-memory single-element forms, hence the extra restriction when is_smem or multimem is set.

Source

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

      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(
            f"f16/bf16 atomic stores require even vector length,"
            f" got {vec_len}"
        )
    i32_vec_len = vec_len * element_bitwidth // 32
    vreg = utils.bitcast(vreg, ir.VectorType.get(
        (i32_vec_len,), i32,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use atomic='add' for smem/multimem f16/bf16 stores
Place the min/max buffer in global memory instead of shared memory
Accumulate min/max in f32 or i32 and narrow afterwards

Example fix

# before
fa16.store_tiled_async(smem_ref, atomic='min')
# after
fa16.store_tiled_async(smem_ref, atomic='add')  # or use a gmem ref for min/max

Defensive patterns

Strategy: validation

Validate before calling

from jax._src.lib import ir
if isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)) and (is_smem or multimem):
    assert atomic == 'add', 'smem/multimem half atomics support add only'

Type guard

from jax._src.lib import ir

def half_smem_atomic_ok(fa, atomic, is_smem, multimem) -> bool:
    if not isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)):
        return True
    return atomic == 'add' or not (is_smem or multimem)

Try / catch

try:
    fa.store_tiled_async(ref, atomic=atomic)
except NotImplementedError:
    if atomic in ('min', 'max'):
        fa.store_tiled_async(gmem_ref, atomic=atomic)  # relocate buffer

Prevention

Keep min/max half-precision buffers in global memory
Use add-only accumulation in shared-memory reduction buffers

When it happens

Trigger: Calling store_tiled_async on an f16/bf16 array with atomic='min' or 'max' while the reference targets shared memory (is_smem) or a multimem descriptor, or when the packed x2 path is required.

Common situations: Using min/max accumulation into a shared-memory reduction buffer at half precision inside a persistent kernel; works when the buffer is in global memory, so the smem restriction surprises developers.

Related errors

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