{"record":{"id":"b611ad9c606355f5","repo":"jax-ml/jax","slug":"wgmma-accumulators-can-only-be-allocated-non-colle","errorCode":null,"errorMessage":"WGMMA accumulators can only be allocated non-collectively. Hint: remove collective_axes from run_scoped. If other allocations are performed as well, split the run_scoped into two.","messagePattern":"WGMMA accumulators can only be allocated non-collectively\\. Hint: remove collective_axes from run_scoped\\. If other allocations are performed as well, split the run_scoped into two\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3587,"sourceCode":"  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. \"\n              f\"Expected signed integer, but got {aval.dtype}.\"\n          )\n\n        dtype = mlir.dtype_to_ir_type(aval.dtype)\n        if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n          input_refs.append(\n              mgpu.WGMMAAccumulator.zero(\n                  *aval.shape, dtype=dtype, is_signed=is_signed\n              )\n          )","sourceCodeStart":3569,"sourceCodeEnd":3605,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3569-L3605","documentation":"WGMMA accumulator buffers (WGMMAAbstractAccumulatorRef) must be allocated per-warpgroup without collectivity; passing collective_axes to a run_scoped that allocates a WGMMA accumulator raises ValueError with a hint to split the allocation.","triggerScenarios":"A single pl.run_scoped(body, collective_axes=(wg,)) whose body allocates both a WGMMA accumulator and other shared buffers.","commonSituations":"Refactoring a matmul kernel so all scratch lives in one scoped allocation, accidentally including the WGMMA accumulator.","solutions":["Split into two run_scoped calls: one non-collective for the WGMMA accumulator, one collective for shared buffers","Remove collective_axes from the run_scoped that allocates the accumulator","Keep accumulator allocation in the warpgroup-private scope by construction"],"exampleFix":"// before\npl.run_scoped(body_with_acc_and_scratch, collective_axes=(wg,))\n// after\npl.run_scoped(alloc_acc_body)\npl.run_scoped(alloc_scratch_body, collective_axes=(wg,))","handlingStrategy":"validation","validationCode":"if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):\n    assert not collective_axes, 'WGMMA accumulators: no collective_axes'","typeGuard":"def is_wgmma_acc(aval) -> bool:\n    return isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef)","tryCatchPattern":null,"preventionTips":["Keep accumulator allocations in dedicated run_scoped calls","Never mix accumulators and shared scratch in one collective scope"],"tags":["jax","pallas","wgmma","run-scoped","collective-axes"],"backgroundTag":"invalid-collective-axes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}