{"record":{"id":"4dfe5d0940b721ee","repo":"jax-ml/jax","slug":"only-thread-collective-allocations-are-supported-i-4dfe5d","errorCode":null,"errorMessage":"Only thread-collective allocations are supported in run_scoped.","messagePattern":"Only thread-collective allocations are supported in run_scoped\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3577,"sourceCode":"def _run_scoped_lowering_rule(\n    ctx: LoweringRuleContext,\n    *consts,\n    jaxpr: jax_core.Jaxpr,\n    collective_axes,\n    **_,\n):\n  if pallas_core.poison_buffers_enabled():\n    raise NotImplementedError(\"Buffer poisoning is not supported on GPU yet.\")\n  input_refs = []\n  should_discharge = []\n  wg_axis = ctx.module_ctx.axis_names.wg\n  is_multithreaded = wg_axis is not None\n  is_thread_collective = is_multithreaded and collective_axes == (wg_axis,)\n  # Make sure everyone has exited previous scoped allocations. Note that we\n  # don't synchronize when we exit the allocation, but only when we might want\n  # to reuse its memory again.\n  if collective_axes and collective_axes != (wg_axis,):\n    raise ValueError(\n        \"Only thread-collective allocations are supported in run_scoped.\"\n    )\n  if is_multithreaded and is_thread_collective:\n    gpu_dialect.barrier()\n  with contextlib.ExitStack() as alloc_stack:\n    for v in jaxpr.invars:\n      aval = cast(ShapedAbstractValue, v.aval)\n      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):\n        if collective_axes:\n          raise ValueError(\n              \"WGMMA accumulators can only be allocated non-collectively. Hint:\"\n              \" remove collective_axes from run_scoped. If other allocations\"\n              \" are performed as well, split the run_scoped into two.\"\n          )\n        is_signed = mgpu_utils.is_signed(aval.dtype)\n        if is_signed is not None and not is_signed:\n          raise ValueError(\n              \"Invalid WGMMA accumulator dtype for s8/i8 WGMMA. \"","sourceCodeStart":3559,"sourceCodeEnd":3595,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3559-L3595","documentation":"In multithreaded (warp-specialized) kernels, run_scoped allocations may only be collective over exactly the warpgroup axis. Passing collective_axes that don't equal (wg_axis,) — e.g. a different axis or multiple axes — raises ValueError.","triggerScenarios":"Calling pl.run_scoped(body, collective_axes=('other',)) or with extra axes in a kernel where the wg axis exists, so collective_axes != (wg_axis,).","commonSituations":"Adding collective allocations to warp-specialized kernels and naming a non-wg mesh axis; mixing cluster and warpgroup axes in one run_scoped.","solutions":["Set collective_axes to exactly the warpgroup axis name (ctx.module_ctx.axis_names.wg) or drop it","Split nested run_scoped calls so each is either wg-collective or non-collective","Check what axis_names.wg is for your kernel and align names"],"exampleFix":"// before\npl.run_scoped(body, collective_axes=('cluster_row',))\n// after\npl.run_scoped(body, collective_axes=wg_axis_name)","handlingStrategy":"validation","validationCode":"assert collective_axes in ((), (wg_axis,)), 'run_scoped axes must be () or (wg,)'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Centralize run_scoped helpers that inject the wg axis","Test warp-specialized kernels on every change to allocations"],"tags":["jax","pallas","run-scoped","collective-axes","warp-specialization"],"backgroundTag":"invalid-collective-axes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}