jax-ml/jax · error · ValueError

Only single-thread arrival is supported on pre-Hopper hardwa

Error message

Only single-thread arrival is supported on pre-Hopper hardware

What it means

On architectures before Hopper (sm < 90), the inline PTX mbarrier.arrive path cannot carry an arrival_count operand, so only arrival_count == 1 is emitted. Requesting multi-thread arrival on pre-Hopper GPUs raises ValueError.

Source

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

        # We need to synchronize the threads after `::before_thread_sync`, as
        # not all threads arrive on the barrier.
        if scope == ThreadSubset.WARPGROUP:
          warpgroup_barrier()
        elif scope == ThreadSubset.WARP:
          warp_barrier()
        else:
          raise ValueError(f"Unsupported scope: {scope}")

    ptx_scope = self._ptx_scope
    if can_complete or ptx_scope != "cta":
      pred_ptx = pred_constraint = ""
      if predicate is not None:
        pred_ptx = "@$2"
        pred_constraint = ",b"
      count_ptx = f", {arrival_count}"
      if get_arch().major < 9:
        if arrival_count != 1:
          raise ValueError(
              "Only single-thread arrival is supported on pre-Hopper hardware"
          )
        count_ptx = ""
      llvm.inline_asm(
          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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make arrival_count configurable and set 1 on pre-Hopper (each thread arrives individually)
  2. Gate the pipeline on jax.experimental.mosaic.gpu.utils.get_arch().major >= 9 and fall back to a single-producer scheme
  3. Have one lane arrive on behalf of the group using a predicate on pre-Hopper hardware

Example fix

# before
barrier.arrive(arrival_count=threads_per_cta)
# after
from jax.experimental.mosaic.gpu.utils import get_arch
count = threads_per_cta if get_arch().major >= 9 else 1
barrier.arrive(arrival_count=count)
Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental.mosaic.gpu.utils import get_arch
assert get_arch().major >= 9 or arrival_count == 1

Prevention

When it happens

Trigger: barrier.arrive(arrival_count=4) on Ampere or earlier (get_arch().major < 9); porting a Hopper pipeline (e.g. TMA warp-specialized code) to A100/sm80.

Common situations: Developing on H100 and running on older GPUs in CI; defaulting arrival_count to warp size for producer arrival; hardcoding multi-thread arrivals copied from Blackwell examples.

Related errors


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