jax-ml/jax · error · ValueError

The cp.async implementation only supports…

Error message

The cp.async implementation only supports oob_mode=OOBFillMode.PROMISE_IN_BOUNDS

What it means

Fires in the gmem-to-smem copy lowering on pre-Hopper GPUs (compute capability major < 9), where TMA is unavailable and the copy is implemented with cp.async. cp.async performs no out-of-bounds handling, so the copy's oob_mode must be OOBFillMode.PROMISE_IN_BOUNDS; requesting bounds-checked fills on this fallback path fails.

Solutions

1. On pre-Hopper GPUs, use oob_mode=OOBFillMode.PROMISE_IN_BOUNDS and guarantee in-bounds copies (e.g. pad the tensor or mask block iteration)
2. Gate the kernel so the ZEROS path only runs on Hopper+ (TMA)
3. Handle boundary blocks manually with predicated loads on older hardware

Example fix

# before
copy_gmem_to_smem(src, smem, oob_mode=OOBFillMode.ZEROS)
# after
mode = OOBFillMode.ZEROS if mgpu.utils.get_arch().major >= 9 else OOBFillMode.PROMISE_IN_BOUNDS
copy_gmem_to_smem(src, smem, oob_mode=mode)

from jax._src.pallas.mosaic_gpu import mgpu
from jax._src.pallas.mosaic_gpu.primitives import OOBFillMode
if mgpu.utils.get_arch().major < 9:
    oob_mode = OOBFillMode.PROMISE_IN_BOUNDS
    # ensure blocks stay in bounds: pad tensors or mask iteration

When it happens

Trigger: Calling copy_gmem_to_smem with oob_mode=OOBFillMode.ZEROS (or any mode other than PROMISE_IN_BOUNDS) on a GPU with compute capability < 9.

Common situations: Defaulting to ZEROS fill for boundary blocks (common in flash-attention style kernels) and running on Ampere or older; writing portable kernels without arch-dependent OOB handling.

Related errors

• copy_gmem_to_smem with a barrier is only supported Hopper…
• Only the TMA implementation supports collective copies
• Only the TMA implementation supports leader_tracked copies
• The cp.async implementation does not support user-defined…
• Out-of-bounds read of

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

Appendix: source

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

  # 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

    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)

View on GitHub (pinned to 1e1c6a8fc0)