jax-ml/jax · error · ValueError

Collective allocations along cluster axes are not supported.

Error message

Collective allocations along cluster axes are not supported.

What it means

In Mosaic GPU interpret mode, scoped memory allocations (run_scoped) may only be collective along the thread axis, not along cluster axes of the warp mesh. The interpreter simulates each thread sequentially and cannot model allocations shared across cluster dimensions, so it raises this ValueError when any cluster axis is marked collective.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py:105

def _raise_if_unsupported_collective_axes(
    mesh: mosaic_gpu_core.Mesh | None,
    is_collective_by_thread_cluster_axis: tuple[bool, ...],
):
  if not mesh or not mesh.thread_name:
    if any(is_collective_by_thread_cluster_axis):
      raise ValueError(
          "Requesting collective allocations, but no explicit thread axis"
          " specified."
      )
  else:
    # Note that the leading entries in `is_collective_by_thread__cluster_axis`
    # correspond to the cluster axes, while the last entry corresponds to the
    # thread axis within a block.
    *is_collective_by_cluster_axis, is_thread_axis_collective = (
        is_collective_by_thread_cluster_axis
    )
    if any(is_collective_by_cluster_axis):
      raise ValueError(
          "Collective allocations along cluster axes are not supported."
      )
    if not is_thread_axis_collective:
      raise ValueError(
          "Scoped allocation must have the thread axis in its collective axes."
      )


# TODO(nrink): Try unifying this function with `_extract_barrier_slice_base`
# from `jax._src.pallas.mosaic_gpu.primitives`.
def _get_index_for_barrier_allocation_key(
    transforms_treedef, transforms_leaves,
) -> indexing.DimIndexer | None:
  # TODO(nrink): The working out of `transforms` and the returned index below
  # may need tidying up. Specifically, GPU interpret mode should correctly
  # support legal ways to index into barriers. (Here, 'legal' is to be read as
  # 'allowed by the Pallas GPU semantics'.)
  if transforms_treedef is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restrict collective_axes in run_scoped to only the thread axis name of your WarpMesh
2. Redesign the kernel to allocate per-warp and synchronize with barriers instead of cluster-wide scoped allocations
3. Run the kernel on real hardware compilation instead of interpret mode if cluster-collective allocation is required

Example fix

// before
run_scoped(alloc, collective_axes=('row', 'col'))  # mesh cluster axes
// after
run_scoped(alloc, collective_axes=('thread',))    # only thread axis

mesh = kernel_launch_mesh
collective = set(collective_axes or ())
cluster_axes = set(mesh.axis_names[:-1]) if mesh else set()
assert not (collective & cluster_axes), 'cluster-axis collective allocations unsupported'

When it happens

Trigger: Calling run_scoped with a WarpMesh whose collective_axes include one or more cluster axes (i.e. anything other than the single thread axis), while interpreting an MGPU kernel.

Common situations: Porting TPU-style collective SMEM allocations to Mosaic GPU kernels; using collective_axes computed from a multi-dimension mesh without filtering to the thread axis.

Related errors

• `thread_id` must be zero when allocating a buffer for all th
• `block_id` must be zero when allocating a buffer for all thr
• Out-of-bounds read of {allocation_key}: reading [{read_range
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• masked load_p

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