jax-ml/jax · error · ValueError

copy_gmem_to_smem with a barrier is only supported Hopper an

Error message

copy_gmem_to_smem with a barrier is only supported Hopper and newer GPUs, which use the TMA implementation

What it means

On GPUs older than Hopper (compute capability < 9, is_cp_async=True), the copy falls back to cp.async, which has no mbarrier support. Passing a barrier to copy_gmem_to_smem on such hardware raises ValueError because only the Hopper+ TMA implementation supports barriers.

Source

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

  if is_leader_tracked_copy:
    # Leader receives the completion messages from both CTAs.
    bytes *= 2
    if len(collective) != 1:
      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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard barrier usage on compute capability: only pass a barrier when mgpu.utils.get_arch().major >= 9
2. Use a non-barrier copy plus explicit synchronization on pre-Hopper GPUs
3. Target Hopper (H100/H200/B100) hardware for barrier-based pipelines

Example fix

# before
copy_gmem_to_smem(src, smem, barrier=bar)
# after
use_tma = mgpu.utils.get_arch().major >= 9
copy_gmem_to_smem(src, smem, barrier=bar if use_tma else None)

from jax._src.pallas.mosaic_gpu import mgpu
use_tma = mgpu.utils.get_arch().major >= 9
assert use_tma or barrier is None, 'barrier copies need Hopper+'

When it happens

Trigger: Running copy_gmem_to_smem(..., barrier=...) on Ampere (SM80) or older GPUs where the TMA path is unavailable.

Common situations: Developing on/for older GPUs (A100, V100) or multi-GPU fleets with mixed generations; code written against Hopper TMA tutorials being run elsewhere.

Related errors

• Only the TMA implementation supports collective copies
• 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
• arrive_expect_tx is only supported on Hopper+ hardware

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