jax-ml/jax · error · ValueError

wait_gmem_to_smem is only supported on pre-Hopper GPUs, whic

Error message

wait_gmem_to_smem is only supported on pre-Hopper GPUs, which use cp.async for GMEM->SMEM copies.

What it means

wait_gmem_to_smem waits for completion of cp.async copies, which are only used on pre-Hopper (SM < 9) GPUs. On Hopper+ the GMEM->SMEM copies go through TMA with barriers, so there is no cp.async group to wait on and the call is invalid.

Source

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


@wait_gmem_to_smem_p.def_effectful_abstract_eval
def _wait_gmem_to_smem_abstract_eval(n):
  del n  # Unused.
  return (), {gpu_core._memory_effect}


@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, *gpu_core.LANExWARP_SEMANTICS)
@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(
    wait_gmem_to_smem_p, *gpu_core.WGxWARP_SEMANTICS)
def _wait_gmem_to_smem_lowering(ctx: lowering.LoweringRuleContext, n):
  if mgpu.utils.get_arch().major >= 9:
    raise ValueError(
        "wait_gmem_to_smem is only supported on pre-Hopper GPUs, which use"
        " cp.async for GMEM->SMEM copies."
    )
  ctx.launch_ctx.await_cp_async_copy(allow_groups=n)
  return ()


def wait_gmem_to_smem(n: int) -> None:
  """Waits until at most ``n`` ``cp.async`` GMEM->SMEM copies are in flight.

  .. note:: This waiting mechanism is only supported on pre-Hopper GPUs, which
            use the ``cp.async`` implementation of
            :func:`jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem`.

  Args:
    n: The maximum number of copies allowed to remain in flight.
  """
  wait_gmem_to_smem_p.bind(n)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Replace copy-without-barrier + wait_gmem_to_smem with an explicit barrier: pass barrier=barrier to copy_gmem_to_smem and call barrier.wait().
Use copy_in / higher-level helpers which select the right synchronization per architecture.

Example fix

# before
copy_gmem_to_smem(src_ref, dst_ref)
wait_gmem_to_smem()
# 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()

Defensive patterns

Strategy: fallback

Validate before calling

from jax._src.pallas.mosaic_gpu import mgpu
pre_hopper = mgpu.utils.get_arch().major < 9
if pre_hopper:
    copy_gmem_to_smem(src, dst); wait_gmem_to_smem()
else:
    barrier = alloc_barrier(); copy_gmem_to_smem(src, dst, barrier=barrier); barrier.wait()

Prevention

Use barriers unconditionally; they also work on pre-Hopper.
Avoid cp.async-specific APIs in kernels targeted at multiple GPU generations.

When it happens

Trigger: Calling wait_gmem_to_smem in a kernel running on a Hopper (SM90) or newer GPU.

Common situations: A100-era kernels with explicit cp.async waits (copy without barrier then wait_gmem_to_smem) run on H100/Blackwell after a cluster migration.

Related errors

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