What causes "Barriers are required for TMA GMEM -> SMEM copies" in jax-ml/jax?

Pass a barrier (e.g. mgpu.Barrier/await barrier object) to async_copy for TMA GMEM -> SMEM loads, and wait on it before reading SMEM.

jax-ml/jax · error · ValueError

Barriers are required for TMA GMEM -> SMEM copies

Error message

Barriers are required for TMA GMEM -> SMEM copies

What it means

TMA loads from GMEM into SMEM are asynchronous hardware operations whose completion must be tracked by an mbarrier. Mosaic therefore requires a barrier for GMEM -> SMEM TMA copies; omitting the barrier argument raises this ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:1289

      collective = ()
    if not isinstance(gmem_transform, tuple):
      gmem_transform = (gmem_transform,)
    if not isinstance(gmem_slice, tuple):
      gmem_slice = (gmem_slice,)

    if reduction_op is not None:
      if implementation != AsyncCopyImplementation.TMA:
        raise ValueError("Only the TMA implementation supports reductions")
      if not _is_tma_reduction_op_supported(reduction_op, element_type):
        raise ValueError(
            f"Reduction op {reduction_op} not supported by the TMA"
            f" implementation for element type {element_type}"
        )

    if src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty):
      gmem_ref, smem_ref = src_ref, dst_ref
      if implementation == AsyncCopyImplementation.TMA and barrier is None:
        raise ValueError("Barriers are required for TMA GMEM -> SMEM copies")
      if arrive is None:
        arrive = True  # Arrive by default
    elif utils.is_smem_ref(src_ref_ty) and dst_ref_ty.memory_space is None:
      gmem_ref, smem_ref = dst_ref, src_ref
      if barrier is not None:
        raise ValueError("Barriers are unsupported for SMEM -> GMEM copies")
      if arrive is None:
        arrive = True  # Commit this copy to the async group by default
    else:
      raise ValueError("Only SMEM <-> GMEM copies supported")

    if collective and gmem_ref is dst_ref:
      raise ValueError("Only GMEM -> SMEM copies can be collective")

    (
        slice_shape,
        untransformed_slice_shape,
        dyn_base_indices,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass a barrier (e.g. mgpu.Barrier/await barrier object) to async_copy for TMA GMEM -> SMEM loads, and wait on it before reading SMEM.
Use arrive=True (the default) so the TMA transaction signals the barrier on completion.
If you cannot provide a barrier, use a non-TMA implementation for this copy.

Example fix

// before
ctx.async_copy(gmem_ref, smem_ref, ..., implementation=mgpu.AsyncCopyImplementation.TMA)
// after
bar = mgpu.Barrier(mgpu.MemRef(() , jnp.uint32), 1)
ctx.async_copy(gmem_ref, smem_ref, ..., barrier=bar, implementation=mgpu.AsyncCopyImplementation.TMA)
bar.await_value(1)  # before consuming smem_ref

Defensive patterns

Strategy: validation

Validate before calling

is_load = src_ref_ty.memory_space is None and utils.is_smem_ref(dst_ref_ty)
if is_load and implementation == mgpu.AsyncCopyImplementation.TMA:
    assert barrier is not None, 'TMA GMEM->SMEM copies require barrier='

Try / catch

try:
    ctx.async_copy(gmem_ref, smem_ref, implementation=mgpu.AsyncCopyImplementation.TMA)
except ValueError as e:
    if 'Barriers are required' in str(e):
        bar = mgpu.Barrier(...)
        ctx.async_copy(gmem_ref, smem_ref, barrier=bar,
                       implementation=mgpu.AsyncCopyImplementation.TMA)
    else:
        raise

Prevention

Always create an mbarrier alongside SMEM buffers used for TMA loads.
Wait on the barrier before reading SMEM to avoid data races.
Use code-generation helpers that always thread barriers through async_copy calls.

When it happens

Trigger: Calling async_copy(src_gmem, dst_smem, implementation=AsyncCopyImplementation.TMA) without passing barrier=..., so the load completion cannot be synchronized.

Common situations: Migrating a working LDGSTS-based copy to TMA and forgetting the mbarrier; new Mosaic kernels copied from examples that omit barrier setup; consuming SMEM data immediately after the copy and hitting races instead.

Related errors

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