jax-ml/jax · error · ValueError

copy_gmem_to_smem without a barrier is only supported on pre

Error message

copy_gmem_to_smem without a barrier is only supported on pre-Hopper GPUs, which use the cp.async implementation

What it means

Mosaic GPU's copy_gmem_to_smem lowering requires an explicit barrier when running on Hopper+ (SM90+) GPUs, because on those architectures the copy is implemented with TMA (async_load) which needs a barrier to signal completion. Only pre-Hopper GPUs use the cp.async path where barrier=None is allowed. The check fires when barrier is None under Warpgroup lowering semantics on TMA-capable hardware.

Source

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

          "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:
      raise ValueError(
          "copy_gmem_to_smem without a barrier is only supported on pre-Hopper"
          " GPUs, which use the cp.async implementation"
      )

  i32 = ir.IntegerType.get_signless(32)
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    if (
        ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warpgroup
        and ctx.module_ctx.auto_barriers
    ):
      mgpu.warpgroup_barrier()  # Make sure all reads have completed.

    if not is_cp_async:
      assert barrier is not None
      if bytes % WARPGROUP_SIZE:
        raise NotImplementedError(
            "Only copies transferring a number of bytes divisible by the"
            f" warpgroup size are supported. Got {bytes=} but warpgroup size is"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a mgpu.SMEM_ALLOCATOR Barrier to copy_gmem_to_smem, e.g. allocate a barrier and pass barrier=barrier, then wait_gmem_to_smem/barrier.wait after the copy on Hopper.
2. If you truly want the no-barrier cp.async path, target a pre-Hopper GPU (set the Mosaic GPU arch/environment to SM80).
3. Upgrade kernel code to the TMA-oriented API (copy_in / device topology helpers) which manages barriers automatically.

Example fix

# before
copy_gmem_to_smem(src_ref, dst_ref)
# after
barrier = mgpu.SMEM_ALLOCATOR.get_buffer((1,), mgpu.BarrierType, barrier_init=1)
copy_gmem_to_smem(src_ref, dst_ref, barrier=barrier)
barrier.wait()

import jax
from jax._src.pallas.mosaic_gpu import mgpu
arch = mgpu.utils.get_arch()
needs_barrier = arch.major >= 9
if needs_barrier:
    barrier = mgpu.SMEM_ALLOCATOR.get_buffer((1,), mgpu.BarrierType, barrier_init=1)

When it happens

Trigger: Calling copy_gmem_to_smem (directly or via copy_in) without passing a barrier while running on a Hopper (SM90) or newer GPU, i.e. when the lowering does not take the cp.async (is_cp_async) path.

Common situations: Kernels written for A100 that ran copy_gmem_to_smem(src, dst) with no barrier are run unchanged on H100; or code that conditionally omits the barrier based on an outdated device check.

Related errors

• copy_gmem_to_smem with a barrier is only supported Hopper an
• wait_gmem_to_smem is only supported on pre-Hopper GPUs, whic
• arrive_expect_tx is only supported on Hopper+ hardware
• complete_tx is only supported on Hopper+ hardware
• Unsupported barrier type: {type(barrier)}

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