{"record":{"id":"1edfa2cfd9132e0f","repo":"jax-ml/jax","slug":"complete-tx-is-only-supported-on-hopper-hardware","errorCode":null,"errorMessage":"complete_tx is only supported on Hopper+ hardware","messagePattern":"complete_tx is only supported on Hopper\\+ hardware","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":1242,"sourceCode":"  def arrive_expect_tx(\n      self, bytes: int | ir.Value, predicate: ir.Value | None = None\n  ):\n    if get_arch().major < 9:\n      raise NotImplementedError(\"arrive_expect_tx is only supported on Hopper+ hardware\")\n    if isinstance(bytes, int):\n      bytes = c(bytes, ir.IntegerType.get_signless(32))\n    elif isinstance(bytes.type, ir.IndexType):\n      i32 = ir.IntegerType.get_signless(32)\n      bytes = arith.index_cast(i32, bytes)\n    nvvm.mbarrier_arrive_expect_tx(\n        self.get_ptr(), bytes, predicate=predicate, scope=self._nvvm_scope\n    )\n\n  def complete_tx(\n      self, bytes: int | ir.Value, predicate: ir.Value | None = None\n  ):\n    if get_arch().major < 9:\n      raise NotImplementedError(\"complete_tx is only supported on Hopper+ hardware\")\n    if isinstance(bytes, int):\n      bytes = c(bytes, ir.IntegerType.get_signless(32))\n    elif isinstance(bytes.type, ir.IndexType):\n      i32 = ir.IntegerType.get_signless(32)\n      bytes = arith.index_cast(i32, bytes)\n\n    pred_ptx = pred_constraint = \"\"\n    if predicate is not None:\n      pred_ptx = \"@$2\"\n      pred_constraint = \",b\"\n\n    llvm.inline_asm(\n        ir.Type.parse(\"!llvm.void\"),\n        [self.get_ptr(), bytes]\n        + ([predicate] if predicate is not None else []),\n        f\"{pred_ptx} mbarrier.complete_tx.shared::{self._ptx_scope}.b64 [$0], $1;\",\n        \"l,r\" + pred_constraint,\n        has_side_effects=True,","sourceCodeStart":1224,"sourceCodeEnd":1260,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L1224-L1260","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nbarrier.complete_tx(bytes=nbytes)\n# after\nfrom jax.experimental.mosaic.gpu.utils import get_arch\nif get_arch().major >= 9:\n    barrier.complete_tx(bytes=nbytes)\nelse:\n    barrier.arrive()  # fallback completion","handlingStrategy":"fallback","validationCode":"from jax.experimental.mosaic.gpu.utils import get_arch\nif get_arch().major < 9:\n    barrier.arrive()  # skip TX accounting on pre-Hopper\nelse:\n    barrier.complete_tx(bytes=nbytes)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Gate the whole TMA pipeline on sm90+, not just individual calls","Keep arrive_expect_tx and complete_tx behind the same arch check"],"tags":["jax","mosaic-gpu","barrier","tma","gpu-architecture","hopper"],"backgroundTag":"gpu-arch-feature-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}