jax-ml/jax · error · NotImplementedError

The cp.async implementation does not support user-defined pr

Error message

The cp.async implementation does not support user-defined predicates

What it means

The cp.async fallback does not support user-supplied predicates (has_user_predicate); on pre-Hopper GPUs passing a predicate to copy_gmem_to_smem raises NotImplementedError. Predicated copies are only available via TMA on Hopper+.

Source

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

      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)
  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
    ):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the predicate on pre-Hopper GPUs and instead mask/skip the copy at the kernel level (e.g. conditionally execute the copy op)
2. Use predicated element-wise loads (regular pl.load with a mask) on older hardware
3. Gate the predicated fast path on get_arch().major >= 9

Example fix

# before
copy_gmem_to_smem(src, smem, predicate=pred)
# after
if mgpu.utils.get_arch().major >= 9:
  copy_gmem_to_smem(src, smem, predicate=pred)
elif pred:
  copy_gmem_to_smem(src, smem)

from jax._src.pallas.mosaic_gpu import mgpu
if mgpu.utils.get_arch().major < 9:
    if not predicate:  # emulate predicated copy
        skip_copy = True  # caller must skip the copy op entirely

When it happens

Trigger: Calling copy_gmem_to_smem(..., predicate=...) on a GPU with compute capability < 9.

Common situations: Using predicates to skip boundary blocks in pipelined loops; kernels written for TMA being run on A100 or older cards.

Related errors

• copy_gmem_to_smem with a barrier is only supported Hopper an
• Only the TMA implementation supports collective copies
• Only the TMA implementation supports leader_tracked copies
• The cp.async implementation only supports oob_mode=OOBFillMo
• Unsupported dtype: {ref.dtype}

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