jax-ml/jax · error · ValueError

Only the TMA implementation supports collective copies

Error message

Only the TMA implementation supports collective copies

What it means

On pre-Hopper GPUs the cp.async fallback is used, and cp.async has no multicast/collective support — only TMA (Hopper+) implements collective copies. Passing collective_axes on such hardware raises ValueError.

Source

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

      raise ValueError(
          f"Expected exactly one collective axis, got {collective_axes=}"
      )
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Gate collective_axes on get_arch().major >= 9 and fall back to per-device copies
2. Issue separate non-collective copies per device on older hardware
3. Run on Hopper+ hardware for multicast paths

Example fix

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

from jax._src.pallas.mosaic_gpu import mgpu
if mgpu.utils.get_arch().major < 9:
    collective_axes = None  # cp.async can't multicast

When it happens

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

Common situations: Multicast kernel developed for H100 run on A100; CI runners with older GPUs; forgetting hardware gating in a shared kernel file.

Related errors

• copy_gmem_to_smem with a barrier is only supported Hopper an
• Only the TMA implementation supports leader_tracked copies
• The cp.async implementation only supports oob_mode=OOBFillMo
• The cp.async implementation does not support user-defined pr
• collective_axes is not supported in pallas_call. Use plgpu.k

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