jax-ml/jax · error · NotImplementedError

Only thread-collective allocations are supported in multithr

Error message

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}).

What it means

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.

Source

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

        should_discharge.append(True)
        continue
      if (
          isinstance(aval, state_types.AbstractRef)
          and aval.memory_space == gpu_core.GMEM
          and jnp.issubdtype(aval.dtype, pallas_core.semaphore)
      ):
        input_ref = alloc_stack.enter_context(
            ctx.module_ctx.reserve_semaphores(
                aval.shape, collective_axes=collective_axes
            )
        )
        input_refs.append(input_ref)
        should_discharge.append(False)
        continue

      # All other allocations must be made collectively across all threads.
      if is_multithreaded and not is_thread_collective:
        raise NotImplementedError(
            "Only thread-collective allocations are supported in multithreaded"
            " kernels. Hint: add"
            f" collective_axes={ctx.module_ctx.axis_names.wg} to your"
            " run_scoped if you intend all threads to share the same"
            f" allocation (currently collective_axes={collective_axes})."
        )
      if isinstance(aval.dtype, gpu_core.BarrierType):
        barrier = _get_barrier(aval, ctx.estimator_ctx.arrival_multiplier)
        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)
        input_refs.append(alloc_stack.enter_context(barrier_ctx))
        should_discharge.append(False)
        continue
      if isinstance(aval.dtype, gpu_core.ClusterBarrierType):
        barrier = _get_cluster_barrier(aval, ctx.module_ctx.axis_names)
        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)
        input_refs.append(alloc_stack.enter_context(barrier_ctx))
        should_discharge.append(False)
        continue

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add collective_axes=ctx.module_ctx.axis_names.wg to the run_scoped call
2. Split WGMMA accumulator allocation into its own non-collective run_scoped (see related error)
3. Verify is_multithreaded expectations: if you didn't mean warp specialization, drop the wg axis from the mesh

Example fix

// before
pl.run_scoped(body, acc_ref, scratch_ref)
// after
pl.run_scoped(body_acc, acc_ref)
pl.run_scoped(body_scratch, scratch_ref, collective_axes=(wg,))

if is_multithreaded and not collective_axes == (wg,):
    collective_axes = (wg,)  # required for shared allocations

When it happens

Trigger: pl.run_scoped(body) with no collective_axes (or wrong axes) in a warp-specialized kernel where the body allocates SMEM/GMEM scratch buffers.

Common situations: 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.

Related errors

• Only thread-collective allocations are supported in run_scop
• 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/9af6ae3663ad5a4a. Report an issue: GitHub.