jax-ml/jax · error · ValueError

Only the TMA implementation supports leader_tracked copies

Error message

Only the TMA implementation supports leader_tracked copies

What it means

On pre-Hopper GPUs (cp.async path), leader_tracked copies are unsupported — the leader/completion-tracking mechanism relies on TMA mbarrier features. Passing leader_tracked on such hardware raises ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1005

      )
    if math.prod(ctx.launch_ctx.cluster_size) != 2:
      raise NotImplementedError(
          "Partitioned loads only supported for clusters of size 2. Got"
          f" cluster size {ctx.launch_ctx.cluster_size}."
      )

  # TMA is only available on Hopper and newer. On older architectures we fall
  # back to the cp.async implementation.
  if is_cp_async := mgpu.utils.get_arch().major < 9:
    if barrier is not None:
      raise ValueError(
          "copy_gmem_to_smem with a barrier is only supported Hopper and newer"
          " GPUs, which use the TMA implementation"
      )
    if collective_axes is not None:
      raise ValueError("Only the TMA implementation supports collective copies")
    if leader_tracked is not None:
      raise ValueError(
          "Only the TMA implementation supports leader_tracked copies"
      )
    # cp.async does not predicate out-of-bounds accesses, so the caller has to
    # guarantee that the copy stays in bounds.
    if oob_mode != OOBFillMode.PROMISE_IN_BOUNDS:
      raise ValueError(
          "The cp.async implementation only supports "
          "oob_mode=OOBFillMode.PROMISE_IN_BOUNDS"
      )
    if has_user_predicate:
      raise NotImplementedError(
          "The cp.async implementation does not support user-defined predicates"
      )
  else:
    if oob_mode is None:
      oob_mode = OOBFillMode.ZEROS

    if barrier is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Gate leader_tracked on get_arch().major >= 9
2. Use ordinary barriers or warp synchronization on pre-Hopper hardware
3. Restrict leader-tracked kernels to Hopper+ devices

Example fix

# before
copy_gmem_to_smem(src, smem, leader_tracked=lt)
# after
if mgpu.utils.get_arch().major >= 9:
  copy_gmem_to_smem(src, smem, leader_tracked=lt)
else:
  copy_gmem_to_smem(src, smem)

from jax._src.pallas.mosaic_gpu import mgpu
if mgpu.utils.get_arch().major < 9:
    leader_tracked = None  # unsupported on cp.async

When it happens

Trigger: Calling copy_gmem_to_smem(..., leader_tracked=...) on a GPU with compute capability < 9.

Common situations: Porting leader-based synchronization schemes to older GPUs; running the same kernel across heterogeneous clusters.

Related errors

• copy_gmem_to_smem with a barrier is only supported Hopper an
• Only the TMA implementation supports collective copies
• The cp.async implementation only supports oob_mode=OOBFillMo
• The cp.async implementation does not support user-defined pr
• Expected exactly one collective axis, got {collective_axes=}

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