{"record":{"id":"e4e0899a9a339a24","repo":"jax-ml/jax","slug":"can-only-await-on-cta-local-barriers","errorCode":null,"errorMessage":"Can only await on CTA-local barriers","messagePattern":"Can only await on CTA-local barriers","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":1141,"sourceCode":"      with when(wait_complete):\n        nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC)\n    return wait_complete\n\n  def test(\n      self,\n      orders_tensor_core: bool = False,\n      scope: ThreadSubset = ThreadSubset.WARPGROUP,\n  ) -> ir.Value:\n    parities = memref.load(self.phases, [])\n    parity, new_parities = self.update_parities(parities)\n    wait_complete = self.test_parity(parity, orders_tensor_core, scope)\n    with when(wait_complete):\n      memref.store(new_parities, self.phases, [])\n    return wait_complete\n\n  def wait_parity(self, parity, orders_tensor_core: bool = False):\n    if self._ptx_scope != \"cta\":\n      raise ValueError(\"Can only await on CTA-local barriers\")\n    i32 = ir.IntegerType.get_signless(32)\n    parity = arith.extui(i32, parity)\n    if get_arch().major < 9:\n      # TODO(apaszke): consider using a single lane + barrier for waiting\n      i1 = ir.IntegerType.get_signless(1)\n      while_op = scf.WhileOp([], [])\n      before_block = while_op.before.blocks.append()\n      with ir.InsertionPoint.at_block_begin(before_block):\n        wait_complete = nvvm.mbarrier_test_wait(self.get_ptr(), parity)\n        wait_incomplete = arith.xori(wait_complete, c(1, i1))\n        scf.condition(wait_incomplete, [])\n      after_block = while_op.after.blocks.append()\n      with ir.InsertionPoint.at_block_begin(after_block):\n        scf.yield_([])\n    else:\n      ticks = arith.constant(i32, 10000000)\n      nvvm.mbarrier_try_wait_parity(self.get_ptr(), parity, ticks)\n    if orders_tensor_core:","sourceCodeStart":1123,"sourceCodeEnd":1159,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L1123-L1159","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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'","Use barrier.wait() only for CTA-local groups; split cross-CTA sync into separate barriers","Verify barrier._ptx_scope == 'cta' before calling wait_parity"],"exampleFix":"# before\nbar = utils.BarrierRef(..., ptx_scope='cluster')\nbar.wait_parity(0)\n# after\nbar = utils.BarrierRef(..., ptx_scope='cta')\nbar.wait_parity(0)","handlingStrategy":"validation","validationCode":"assert getattr(barrier, '_ptx_scope', 'cta') == 'cta', 'wait_parity requires CTA-local barrier'","typeGuard":"def is_cta_barrier(b) -> bool:\n    return getattr(b, '_ptx_scope', 'cta') == 'cta'","tryCatchPattern":null,"preventionTips":["Construct barriers with ptx_scope='cta' unless you specifically need cluster scope","Keep arrive and wait calls on the same scope"],"tags":["jax","mosaic-gpu","barrier","scope","ptx"],"backgroundTag":"barrier-scope-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}