jax-ml/jax · error · ValueError

cluster= must be at most 3D, got {self}.

Error message

cluster= must be at most 3D, got {self}.

What it means

Raised by WarpSpecialized or launch-config dataclass __post_init__ validation: the `cluster` tuple describing the thread-block cluster shape has more than 3 dimensions. CUDA clusters are limited to at most 3D (x, y, z), so Mosaic rejects longer tuples at config construction time.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1640

    )


_WARPGROUP_AXIS_NAME = object()

@dataclasses.dataclass(frozen=True, kw_only=True)
class Mesh(pallas_core.Mesh):
  grid: Sequence[int] = ()
  grid_names: Sequence[str] = ()
  cluster: Sequence[int] = ()
  cluster_names: Sequence[str] = ()
  # Those are NOT CUDA threads. On Hopper they correspond to warpgroups.
  num_threads: int | None = None
  thread_name: str | None = None
  kernel_name: str | None = None

  def __post_init__(self):
    if len(self.cluster) > 3:
      raise ValueError(f"cluster= must be at most 3D, got {self}.")
    if len(self.grid_names) != len(self.grid):
      raise ValueError(
          f"grid_names must have the same length as grid, got {self}."
      )
    if len(self.cluster_names) != len(self.cluster):
      raise ValueError(
          f"cluster_names must have the same length as cluster, got {self}."
      )
    if (self.thread_name is None) != (self.num_threads is None):
      raise ValueError(
          "num_threads and thread_name must be either both set or both None,"
          f" got {self}"
      )
    max_mosaic_threads = 2048 // 128
    if self.num_threads is not None and self.num_threads > max_mosaic_threads:
      raise ValueError(
          "Requested too many CUDA threads per block. Each Mosaic thread"
          f" corresponds to 128 CUDA threads. At most {max_mosaic_threads}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce cluster= to at most 3 dimensions, e.g. `cluster=(2, 2, 2)`
2. Fold extra parallelism into grid= instead of cluster=

Example fix

// before
cluster=(2, 2, 2, 2)

// after
cluster=(2, 2, 2)  # move the 4th factor into grid

assert len(cluster) <= 3, f'cluster must be <= 3D, got {len(cluster)}'

def is_valid_cluster(cluster: tuple) -> bool:
    return 1 <= len(cluster) <= 3

try:
    cfg = MyConfig(cluster=cluster)
except ValueError as e:
    cluster = cluster[:3]  # or rebalance into grid
    cfg = MyConfig(cluster=cluster)

When it happens

Trigger: Creating a GPU kernel launch/config object (e.g. a WarpSpecialized descriptor or launch metadata) with `cluster=(2,2,2,2)` or any tuple of length > 3.

Common situations: Copying a multi-dimensional `grid` (which may be longer) into `cluster=`; refactoring code where grid and cluster were previously the same tuple; upgrading JAX versions where cluster support tightened validation.

Related errors

• dims and idxs must have the same length
• grid_names must have the same length as grid, got {self}.
• cluster_names must have the same length as cluster, got {sel
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• num_threads and thread_name must be either both set or both

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