jax-ml/jax · error · ValueError

Group sizes {group_sizes.shape=} must match first dimension

Error message

Group sizes {group_sizes.shape=} must match first dimension of {A.shape=}

What it means

In the Pallas GPU grouped-matmul kernel, group_sizes must be a 1-D array whose length equals the number of groups g (the first dimension of the 3-D weights A). A mismatch means the kernel cannot map each group to its row-range in the ragged output.

Source

Thrown at jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py:229

  trans_rhs: bool = False,
  interpret: bool = False,
  compute_dtype: DTypeLike | None = None,
  acc_dtype: DTypeLike | None = np.float32,
  num_warps: int | None = None,
  num_stages: int | None = None,
  chunk_m: int = CHUNK_M,
  out_dtype: DTypeLike | None = None,
) -> Array:
  """Compute grouped matmul on GPU via a Pallas lowering."""

  msg = "This gmm kernel only supports either (m, k) x (g, k, n) -> (m, n) "
  msg += f"or (m, k) x (g, n, k) -> (m, n), but got {x.shape=} {A.shape=}"
  if not (A.ndim == 3 and x.ndim == 2):
    raise ValueError(msg)
  msg = f"Group sizes {group_sizes.shape=} must match first dimension of "
  msg += f"{A.shape=}"
  if not A.shape[:1] == group_sizes.shape:
    raise ValueError(msg)
  n = A.shape[-1] if not trans_rhs else A.shape[-2]
  Ak = A.shape[-2] if not trans_rhs else A.shape[-1]
  assert Ak == x.shape[1], msg
  size = RaggedDotSizes(m=x.shape[0], k=x.shape[1], n=n, g=A.shape[0])

  # normalize the block sizes for GPU
  block_m, block_k, block_n = (
    pl.next_power_of_2(min(b, s))
    for b, s in zip([block_m, block_k, block_n], [size.m, size.k, size.n])
  )
  block_k, block_n = max(block_k, 16), max(block_n, 16)

  A_spec = pl.BlockSpec((size.g, size.k, block_n), lambda i, j: (0, 0, j))
  if trans_rhs:  # transposed spec
    A_spec = pl.BlockSpec((size.g, block_n, size.k), lambda i, j: (0, j, 0))

  group_metadata = _make_gmm_group_metadata(
    group_sizes=group_sizes, m=size.m, chunk_m=chunk_m

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify A.shape[0] == group_sizes.shape[0] and reconcile the expert count on both sides
2. Build group_sizes from the same routing computation that permuted the weights (e.g. cumsum-derived sizes of the sorted token assignment)
3. Add an explicit assert before gmm to fail early with your own message

Example fix

// before
out = gmm(x, A, group_sizes)  # len(group_sizes) != A.shape[0]

// after
assert group_sizes.shape[0] == A.shape[0]
out = gmm(x, A, group_sizes)

assert group_sizes.shape == A.shape[:1], (group_sizes.shape, A.shape)

When it happens

Trigger: Calling gmm with A.shape == (g, k, n) but group_sizes.shape != (g,), e.g. passing per-expert sizes of length num_experts while A packs a different number of groups (or passing a scalar/split array).

Common situations: MoE routing where group_sizes comes from a routing histogram over a different expert count than the weight tensor; off-by-one or including a padding group in A but not in group_sizes.

Related errors

• This gmm kernel only supports either (m, k) x (g, k, n) -> (
• group_offset is not currently supported in the pallas-triton
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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