jax-ml/jax · error · NotImplementedError

f16/bf16 only supports add, min, max atomics, got {atomic}

Error message

f16/bf16 only supports add, min, max atomics, got {atomic}

What it means

PTX provides red/atom for f16/bf16 only for .add, .min and .max (packed f16x2 forms); there are no bitwise or other arithmetic atomic forms for half types, so store_tiled_async validates the requested atomic op against that set.

Source

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

    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(
            f"f16/bf16 atomic stores require even vector length,"
            f" got {vec_len}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch to atomic='add'/'min'/'max' for f16/bf16 data
2. Use an integer buffer (i32, which supports and/or/xor as b32) for bitwise reductions
3. Bitcast only if semantics actually match — generally keep bitwise atomics on integer types

Example fix

# before
fa16.store_tiled_async(ref, atomic='or')
# after
mask_i32 = compute_masks_as_i32()
mask_i32.store_tiled_async(ref_i32, atomic='or')

from jax._src.lib import ir
HALF_ATOMICS = ('add', 'min', 'max')
if isinstance(fa.mlir_dtype, (ir.F16Type, ir.BF16Type)):
    assert atomic in HALF_ATOMICS, f'use one of {HALF_ATOMICS}'

from jax._src.lib import ir

def half_atomic_supported(atomic) -> bool:
    return atomic in ('add', 'min', 'max')

try:
    fa.store_tiled_async(ref, atomic=atomic)
except NotImplementedError:
    if atomic in ('and', 'or', 'xor'):
        raise  # bitwise needs an int buffer; do not silently change semantics

When it happens

Trigger: Calling store_tiled_async on an F16Type/BF16Type array with atomic='and'/'or'/'xor' (or any op outside add/min/max).

Common situations: Generic bitwise-reduction code reused across dtypes; changing an integer mask-accumulation buffer to bfloat16 without adjusting the atomic op.

Related errors

• f16/bf16 not supported for async atomics
• 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/1d2997c836768dcb. Report an issue: GitHub.