jax-ml/jax · error · ValueError

Only the TMA implementation supports reductions

Error message

Only the TMA implementation supports reductions

What it means

async_copy supports reduction operations (e.g. accumulate-on-load via TMA reduction semantics) only through the TMA implementation. Passing reduction_op with implementation != AsyncCopyImplementation.TMA raises this ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:1279

    element_bitwidth = utils.bitwidth(element_type)
    if element_type != dst_ref_ty.element_type:
      raise ValueError(
          f"Expected same element type, got {element_type} and"
          f" {dst_ref_ty.element_type}"
      )

    if isinstance(collective, gpu.Dimension):
      collective = (collective,)
    elif collective is None:
      collective = ()
    if not isinstance(gmem_transform, tuple):
      gmem_transform = (gmem_transform,)
    if not isinstance(gmem_slice, tuple):
      gmem_slice = (gmem_slice,)

    if reduction_op is not None:
      if implementation != AsyncCopyImplementation.TMA:
        raise ValueError("Only the TMA implementation supports reductions")
      if not _is_tma_reduction_op_supported(reduction_op, element_type):
        raise ValueError(
            f"Reduction op {reduction_op} not supported by the TMA"
            f" implementation for element type {element_type}"
        )

    if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty):
      gmem_ref, smem_ref = src_ref, dst_ref
      if implementation == AsyncCopyImplementation.TMA and barrier is None:
        raise ValueError("Barriers are required for TMA GMEM -> SMEM copies")
      if arrive is None:
        arrive = True  # Arrive by default
    elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None:
      gmem_ref, smem_ref = dst_ref, src_ref
      if barrier is not None:
        raise ValueError("Barriers are unsupported for SMEM -> GMEM copies")
      if arrive is None:
        arrive = True  # Commit this copy to the async group by default

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set implementation=AsyncCopyImplementation.TMA when reduction_op is not None.
Or drop reduction_op and perform the reduction manually after the copy (e.g. accumulate in SMEM/registers).
Confirm the target GPU supports TMA reductions for the element type.

Example fix

// before
ctx.async_copy(src, dst, ..., reduction_op=mgpu.ReductionOp.ADD, implementation=mgpu.AsyncCopyImplementation.LDGSTS)
// after
ctx.async_copy(src, dst, ..., reduction_op=mgpu.ReductionOp.ADD, implementation=mgpu.AsyncCopyImplementation.TMA)

Defensive patterns

Strategy: validation

Validate before calling

if reduction_op is not None:
    assert implementation == mgpu.AsyncCopyImplementation.TMA, \
        'reductions require the TMA implementation'

Try / catch

try:
    ctx.async_copy(..., reduction_op=op)
except ValueError as e:
    if 'Only the TMA implementation supports reductions' in str(e):
        ctx.async_copy(...)  # retry without reduction_op, reduce manually
    else:
        raise

Prevention

Gate reduction usage on TMA availability of the target GPU.
Keep implementation explicit in all async_copy calls.
Reduce manually when portability matters more than the fused reduction.

When it happens

Trigger: Calling async_copy(reduction_op=..., implementation=AsyncCopyImplementation.LDGSTS) (or any non-TMA implementation), which has no hardware support for fused reductions.

Common situations: Porting kernels between GPUs: using reductions on loads that relied on TMA on Hopper but running/configuring a CUDA-core or LDGSTS path; leaving a legacy implementation argument set when adding reduction_op.

Related errors

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