jax-ml/jax · error · ValueError

Can only await on CTA-local barriers

Error message

Can only await on CTA-local barriers

What it means

wait_parity uses PTX mbarrier.try_wait.parity which is only awaited correctly in a CTA-local shared-memory scope. If the BarrierRef was constructed with a ptx_scope other than 'cta' (e.g. 'cluster' or 'gpu'), awaiting its parity raises ValueError.

Source

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

      with when(wait_complete):
        nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC)
    return wait_complete

  def test(
      self,
      orders_tensor_core: bool = False,
      scope: ThreadSubset = ThreadSubset.WARPGROUP,
  ) -> ir.Value:
    parities = memref.load(self.phases, [])
    parity, new_parities = self.update_parities(parities)
    wait_complete = self.test_parity(parity, orders_tensor_core, scope)
    with when(wait_complete):
      memref.store(new_parities, self.phases, [])
    return wait_complete

  def wait_parity(self, parity, orders_tensor_core: bool = False):
    if self._ptx_scope != "cta":
      raise ValueError("Can only await on CTA-local barriers")
    i32 = ir.IntegerType.get_signless(32)
    parity = arith.extui(i32, parity)
    if get_arch().major < 9:
      # TODO(apaszke): consider using a single lane + barrier for waiting
      i1 = ir.IntegerType.get_signless(1)
      while_op = scf.WhileOp([], [])
      before_block = while_op.before.blocks.append()
      with ir.InsertionPoint.at_block_begin(before_block):
        wait_complete = nvvm.mbarrier_test_wait(self.get_ptr(), parity)
        wait_incomplete = arith.xori(wait_complete, c(1, i1))
        scf.condition(wait_incomplete, [])
      after_block = while_op.after.blocks.append()
      with ir.InsertionPoint.at_block_begin(after_block):
        scf.yield_([])
    else:
      ticks = arith.constant(i32, 10000000)
      nvvm.mbarrier_try_wait_parity(self.get_ptr(), parity, ticks)
    if orders_tensor_core:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Await cluster-scope barriers with the cluster-scope wait API (e.g. mbarrier.try_wait.parity.cluster) or construct the barrier with ptx_scope='cta'
2. Use barrier.wait() only for CTA-local groups; split cross-CTA sync into separate barriers
3. Verify barrier._ptx_scope == 'cta' before calling wait_parity

Example fix

# before
bar = utils.BarrierRef(..., ptx_scope='cluster')
bar.wait_parity(0)
# after
bar = utils.BarrierRef(..., ptx_scope='cta')
bar.wait_parity(0)

assert getattr(barrier, '_ptx_scope', 'cta') == 'cta', 'wait_parity requires CTA-local barrier'

def is_cta_barrier(b) -> bool:
    return getattr(b, '_ptx_scope', 'cta') == 'cta'

When it happens

Trigger: barrier.wait_parity(p) on a barrier created with a cluster/gpu ptx_scope; passing a non-'cta' scope string when building the barrier group.

Common situations: Experimenting with distributed shared memory / cluster barriers on Hopper+; copy-pasting barrier construction code that parameterizes ptx_scope; mismatch between the scope used to arrive and to wait.

Related errors

• Unsupported scope: {scope}
• Indexing barrier with {transforms} not supported in GPU inte
• Expected an `NDIndexer`, but got {transforms[0]}
• Attempting to operate on barrier without indexing, but `num_
• Expected a single barrier, got a barrier reference with shap

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