jax-ml/jax · error · ValueError

WGMMA accumulators can only be allocated non-collectively. H

Error message

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.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3587

  should_discharge = []
  wg_axis = ctx.module_ctx.axis_names.wg
  is_multithreaded = wg_axis is not None
  is_thread_collective = is_multithreaded and collective_axes == (wg_axis,)
  # Make sure everyone has exited previous scoped allocations. Note that we
  # don't synchronize when we exit the allocation, but only when we might want
  # to reuse its memory again.
  if collective_axes and collective_axes != (wg_axis,):
    raise ValueError(
        "Only thread-collective allocations are supported in run_scoped."
    )
  if is_multithreaded and is_thread_collective:
    gpu_dialect.barrier()
  with contextlib.ExitStack() as alloc_stack:
    for v in jaxpr.invars:
      aval = cast(ShapedAbstractValue, v.aval)
      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):
        if collective_axes:
          raise ValueError(
              "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."
          )
        is_signed = mgpu_utils.is_signed(aval.dtype)
        if is_signed is not None and not is_signed:
          raise ValueError(
              "Invalid WGMMA accumulator dtype for s8/i8 WGMMA. "
              f"Expected signed integer, but got {aval.dtype}."
          )

        dtype = mlir.dtype_to_ir_type(aval.dtype)
        if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
          input_refs.append(
              mgpu.WGMMAAccumulator.zero(
                  *aval.shape, dtype=dtype, is_signed=is_signed
              )
          )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split into two run_scoped calls: one non-collective for the WGMMA accumulator, one collective for shared buffers
2. Remove collective_axes from the run_scoped that allocates the accumulator
3. Keep accumulator allocation in the warpgroup-private scope by construction

Example fix

// before
pl.run_scoped(body_with_acc_and_scratch, collective_axes=(wg,))
// after
pl.run_scoped(alloc_acc_body)
pl.run_scoped(alloc_scratch_body, collective_axes=(wg,))

if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):
    assert not collective_axes, 'WGMMA accumulators: no collective_axes'

def is_wgmma_acc(aval) -> bool:
    return isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef)

When it happens

Trigger: A single pl.run_scoped(body, collective_axes=(wg,)) whose body allocates both a WGMMA accumulator and other shared buffers.

Common situations: Refactoring a matmul kernel so all scratch lives in one scoped allocation, accidentally including the WGMMA accumulator.

Related errors

• Only thread-collective allocations are supported in run_scop
• Only thread-collective allocations are supported in multithr
• collective_axes is not supported in pallas_call. Use plgpu.k
• Preinitialized WGMMAAccumulatorRef only supported in pl.run_
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

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