jax-ml/jax · error · NotImplementedError

complete_tx is only supported on Hopper+ hardware

Error message

complete_tx is only supported on Hopper+ hardware

What it means

complete_tx lowers to nvvm.mbarrier_complete_tx, which finishes the transaction-byte accounting started by arrive_expect_tx; that instruction exists only on Hopper (sm90)+. On older GPUs Mosaic raises NotImplementedError.

Source

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

  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)
      bytes = arith.index_cast(i32, bytes)

    pred_ptx = pred_constraint = ""
    if predicate is not None:
      pred_ptx = "@$2"
      pred_constraint = ",b"

    llvm.inline_asm(
        ir.Type.parse("!llvm.void"),
        [self.get_ptr(), bytes]
        + ([predicate] if predicate is not None else []),
        f"{pred_ptx} mbarrier.complete_tx.shared::{self._ptx_scope}.b64 [$0], $1;",
        "l,r" + pred_constraint,
        has_side_effects=True,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Branch on get_arch().major >= 9 and use a cp.async + arrive fallback on older GPUs
Ensure the kernel runs only on Hopper+ (CUDA_VISIBLE_DEVICES / device placement)
Pair complete_tx strictly with arrive_expect_tx and gate both behind the same arch check

Example fix

# before
barrier.complete_tx(bytes=nbytes)
# after
from jax.experimental.mosaic.gpu.utils import get_arch
if get_arch().major >= 9:
    barrier.complete_tx(bytes=nbytes)
else:
    barrier.arrive()  # fallback completion

Defensive patterns

Strategy: fallback

Validate before calling

from jax.experimental.mosaic.gpu.utils import get_arch
if get_arch().major < 9:
    barrier.arrive()  # skip TX accounting on pre-Hopper
else:
    barrier.complete_tx(bytes=nbytes)

Prevention

Gate the whole TMA pipeline on sm90+, not just individual calls
Keep arrive_expect_tx and complete_tx behind the same arch check

When it happens

Trigger: barrier.complete_tx(bytes=N) anywhere in async_copy/kernel body code when get_arch().major < 9 — typically paired with a preceding arrive_expect_tx on non-Hopper hardware.

Common situations: TMA/async-copy producer-consumer pipelines ported to A100 or earlier; forgetting that the same kernel module is JIT'd for whatever GPU is visible; mixed-GPU test machines.

Related errors

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