jax-ml/jax · error · ValueError

Requested too many CUDA threads per block. Each Mosaic threa

Error message

Requested too many CUDA threads per block. Each Mosaic thread corresponds to 128 CUDA threads. At most {max_mosaic_threads} are supported, got {self}

What it means

Launch-config validation: each Mosaic 'thread' maps to 128 CUDA threads and a CUDA block supports at most 2048 threads, so num_threads may not exceed 2048 // 128 = 16. Requesting more raises this error before launch.

Source

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

  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}"
          f" are supported, got {self}"
      )
    object.__setattr__(self, "grid", tuple(self.grid))
    object.__setattr__(self, "grid_names", tuple(self.grid_names))
    object.__setattr__(self, "cluster", tuple(self.cluster))
    object.__setattr__(self, "cluster_names", tuple(self.cluster_names))

  @property
  def default_memory_space(self) -> MemorySpace:
    return MemorySpace.GMEM

  @property
  def shape(self) -> collections.OrderedDict[object, int]:
    pairs: Iterable[tuple[object, int]]
    if self.num_threads is not None:
      pairs = zip(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reduce num_threads to at most 16
If more parallelism is needed, restructure across warps/cluster instead of more Mosaic threads

Example fix

// before
num_threads=32

// after
num_threads=16  # max: 2048 CUDA threads // 128

Defensive patterns

Strategy: validation

Validate before calling

MAX_MOSAIC_THREADS = 2048 // 128  # 16
assert num_threads is None or num_threads <= MAX_MOSAIC_THREADS

Type guard

def threads_supported(num_threads: int | None) -> bool:
    return num_threads is None or 1 <= num_threads <= 16

Try / catch

try:
    cfg = LaunchConfig(num_threads=n, thread_name='warp')
except ValueError:
    cfg = LaunchConfig(num_threads=min(n, 16), thread_name='warp')

Prevention

Remember each Mosaic thread = 128 CUDA threads; hard cap is 16
Clamp num_threads at config creation
Expose num_threads as a tunable bounded to 16

When it happens

Trigger: Passing num_threads > 16 (e.g. num_threads=32) to the Mosaic GPU kernel launch config while also setting thread_name.

Common situations: Tuning warp-specialized kernels and assuming num_threads maps 1:1 to CUDA warps; porting kernels between Mosaic versions where the thread model changed to 128-CUDA-thread units.

Related errors

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