{"record":{"id":"9af6ae3663ad5a4a","repo":"jax-ml/jax","slug":"only-thread-collective-allocations-are-supported-i-9af6ae","errorCode":null,"errorMessage":"Only thread-collective allocations are supported in multithreaded kernels. Hint: add collective_axes={ctx.module_ctx.axis_names.wg} to your run_scoped if you intend all threads to share the same allocation (currently collective_axes={collective_axes}).","messagePattern":"Only thread-collective allocations are supported in multithreaded kernels\\. Hint: add collective_axes=(.+?) to your run_scoped if you intend all threads to share the same allocation \\(currently collective_axes=(.+?)\\)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3631,"sourceCode":"        should_discharge.append(True)\n        continue\n      if (\n          isinstance(aval, state_types.AbstractRef)\n          and aval.memory_space == gpu_core.GMEM\n          and jnp.issubdtype(aval.dtype, pallas_core.semaphore)\n      ):\n        input_ref = alloc_stack.enter_context(\n            ctx.module_ctx.reserve_semaphores(\n                aval.shape, collective_axes=collective_axes\n            )\n        )\n        input_refs.append(input_ref)\n        should_discharge.append(False)\n        continue\n\n      # All other allocations must be made collectively across all threads.\n      if is_multithreaded and not is_thread_collective:\n        raise NotImplementedError(\n            \"Only thread-collective allocations are supported in multithreaded\"\n            \" kernels. Hint: add\"\n            f\" collective_axes={ctx.module_ctx.axis_names.wg} to your\"\n            \" run_scoped if you intend all threads to share the same\"\n            f\" allocation (currently collective_axes={collective_axes}).\"\n        )\n      if isinstance(aval.dtype, gpu_core.BarrierType):\n        barrier = _get_barrier(aval, ctx.estimator_ctx.arrival_multiplier)\n        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)\n        input_refs.append(alloc_stack.enter_context(barrier_ctx))\n        should_discharge.append(False)\n        continue\n      if isinstance(aval.dtype, gpu_core.ClusterBarrierType):\n        barrier = _get_cluster_barrier(aval, ctx.module_ctx.axis_names)\n        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)\n        input_refs.append(alloc_stack.enter_context(barrier_ctx))\n        should_discharge.append(False)\n        continue","sourceCodeStart":3613,"sourceCodeEnd":3649,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3613-L3649","documentation":"In multithreaded kernels (wg axis present), every non-WGMMA, non-barrier allocation inside run_scoped must be collective over the warpgroup axis; allocating per-thread raises NotImplementedError with a hint to add collective_axes=wg.","triggerScenarios":"pl.run_scoped(body) with no collective_axes (or wrong axes) in a warp-specialized kernel where the body allocates SMEM/GMEM scratch buffers.","commonSituations":"Taking a single-threaded kernel and enabling warp specialization without updating its run_scoped calls; forgetting that multithreaded kernels require shared allocations to avoid per-thread duplication/races.","solutions":["Add collective_axes=ctx.module_ctx.axis_names.wg to the run_scoped call","Split WGMMA accumulator allocation into its own non-collective run_scoped (see related error)","Verify is_multithreaded expectations: if you didn't mean warp specialization, drop the wg axis from the mesh"],"exampleFix":"// before\npl.run_scoped(body, acc_ref, scratch_ref)\n// after\npl.run_scoped(body_acc, acc_ref)\npl.run_scoped(body_scratch, scratch_ref, collective_axes=(wg,))","handlingStrategy":"validation","validationCode":"if is_multithreaded and not collective_axes == (wg,):\n    collective_axes = (wg,)  # required for shared allocations","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass collective_axes=(wg,) for shared scratch in warp-specialized kernels","Split WGMMA allocations out","Codify the pattern in a helper"],"tags":["jax","pallas","run-scoped","warp-specialization","collective-axes"],"backgroundTag":"invalid-collective-axes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}