jax-ml/jax · error · NotImplementedError

arrive_expect_tx is only supported on Hopper+ hardware

Error message

arrive_expect_tx is only supported on Hopper+ hardware

What it means

arrive_expect_tx uses the Hopper+ PTX instruction mbarrier.arrive.expect_tx, which reserves transaction-byte counting for TMA-style async transfers. On architectures with compute capability < 9.0 it raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1228

          ir.IntegerType.get_signless(64),
          [self.get_ptr()] + ([predicate] if predicate is not None else []),
          f"{pred_ptx} mbarrier.arrive.release.{ptx_scope}.shared::{ptx_scope}.b64 $0, [$1]{count_ptx};",
          "=l,r" + pred_constraint,
          has_side_effects=True,
      )
    else:
      if predicate is not None:
        raise NotImplementedError(
            "Predicate not supported for no-complete arrive"
        )
      count = c(arrival_count, ir.IntegerType.get_signless(32))
      nvvm.mbarrier_arrive_nocomplete(self.get_ptr(), count)

  def arrive_expect_tx(
      self, bytes: int | ir.Value, predicate: ir.Value | None = None
  ):
    if get_arch().major < 9:
      raise NotImplementedError("arrive_expect_tx is only supported on Hopper+ hardware")
    if isinstance(bytes, int):
      bytes = c(bytes, ir.IntegerType.get_signless(32))
    elif isinstance(bytes.type, ir.IndexType):
      i32 = ir.IntegerType.get_signless(32)
      bytes = arith.index_cast(i32, bytes)
    nvvm.mbarrier_arrive_expect_tx(
        self.get_ptr(), bytes, predicate=predicate, scope=self._nvvm_scope
    )

  def complete_tx(
      self, bytes: int | ir.Value, predicate: ir.Value | None = None
  ):
    if get_arch().major < 9:
      raise NotImplementedError("complete_tx is only supported on Hopper+ hardware")
    if isinstance(bytes, int):
      bytes = c(bytes, ir.IntegerType.get_signless(32))
    elif isinstance(bytes.type, ir.IndexType):
      i32 = ir.IntegerType.get_signless(32)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Target only sm90+ (H100/H200/B100) for expect_tx pipelines; guard with get_arch().major >= 9
2. Fall back to cp.async-based transfers with plain arrive/complete on older GPUs
3. Set the platform/arch explicitly in the kernel launch config so get_arch() matches the target device

Example fix

# before
barrier.arrive_expect_tx(bytes=chunk_bytes)
# after
from jax.experimental.mosaic.gpu.utils import get_arch
if get_arch().major >= 9:
    barrier.arrive_expect_tx(bytes=chunk_bytes)
else:
    barrier.arrive()  # legacy cp.async path

from jax.experimental.mosaic.gpu.utils import get_arch
hopper_plus = get_arch().major >= 9

When it happens

Trigger: barrier.arrive_expect_tx(bytes=N) on Ampere/Turing or any GPU where get_arch().major < 9; TMA async-copy pipelines run on pre-Hopper hardware.

Common situations: Running a TMA-based Mosaic kernel written for H100 on A100; CI fleets with mixed GPU generations; relying on cp.async.bulk semantics that only exist on sm90+.

Related errors

• complete_tx is only supported on Hopper+ hardware
• copy_gmem_to_smem with a barrier is only supported Hopper an
• Only single-thread arrival is supported on pre-Hopper hardwa
• copy_gmem_to_smem without a barrier is only supported on pre
• Barriers are required for TMA GMEM -> SMEM copies

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