jax-ml/jax · error · ValueError

Only thread-collective allocations are supported in run_scop

Error message

Only thread-collective allocations are supported in run_scoped.

What it means

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.

Source

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

def _run_scoped_lowering_rule(
    ctx: LoweringRuleContext,
    *consts,
    jaxpr: jax_core.Jaxpr,
    collective_axes,
    **_,
):
  if pallas_core.poison_buffers_enabled():
    raise NotImplementedError("Buffer poisoning is not supported on GPU yet.")
  input_refs = []
  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. "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set collective_axes to exactly the warpgroup axis name (ctx.module_ctx.axis_names.wg) or drop it
2. Split nested run_scoped calls so each is either wg-collective or non-collective
3. Check what axis_names.wg is for your kernel and align names

Example fix

// before
pl.run_scoped(body, collective_axes=('cluster_row',))
// after
pl.run_scoped(body, collective_axes=wg_axis_name)

assert collective_axes in ((), (wg_axis,)), 'run_scoped axes must be () or (wg,)'

When it happens

Trigger: 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,).

Common situations: Adding collective allocations to warp-specialized kernels and naming a non-wg mesh axis; mixing cluster and warpgroup axes in one run_scoped.

Related errors

• Only thread-collective allocations are supported in multithr
• WGMMA accumulators can only be allocated non-collectively. H
• collective_axes is not supported in pallas_call. Use plgpu.k
• num_threads and thread_name must be either both set or both
• `thread_id` must be zero when allocating a buffer for all th

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