jax-ml/jax · error · ValueError

Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},) whe

Error message

Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where `max_num_seqs` is `page_indices.shape[0]`.

What it means

static_validate_inputs requires cu_q_lens (cumulative query lengths, the ragged-layout offset array) to have exactly max_num_seqs + 1 entries, where max_num_seqs = page_indices.shape[0]. It is the standard prefix-sum layout: entry i is the start of sequence i and the last entry is the total query count.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py:255

  _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape
  assert num_combined_kv_heads % 2 == 0
  assert isinstance(k_scale, float) or k_scale is None
  assert isinstance(v_scale, float) or v_scale is None
  num_kv_heads = num_combined_kv_heads // 2
  max_num_seqs, pages_per_seq = page_indices.shape
  if num_seqs.shape != (1,):
    raise ValueError(f"{num_seqs.shape=} must be (1,)")
  if head_dim_k != head_dim:
    raise ValueError(
        f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}."
    )
  if kv_lens.shape != (max_num_seqs,):
    raise ValueError(
        f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if cu_q_lens.shape != (max_num_seqs + 1,):
    raise ValueError(
        f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if (
      kv_lens.dtype != jnp.int32
      or page_indices.dtype != jnp.int32
      or cu_q_lens.dtype != jnp.int32
  ):
    raise ValueError(
        "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be"
        f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},"
        f" {cu_q_lens.dtype=}."
    )
  if num_q_heads % num_kv_heads != 0:
    raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}")
  if sliding_window is not None and sliding_window <= 0:
    raise ValueError(f"{sliding_window=} must be positive.")
  if soft_cap is not None and soft_cap == 0.0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Compute cu_q_lens = jnp.concatenate([[0], jnp.cumsum(q_lens)]) so its length is num_seqs + 1
Verify cu_q_lens.shape == (page_indices.shape[0] + 1,) before the call
Do not pass per-sequence lengths where cumulative offsets are expected

Example fix

// before
cu_q_lens = q_lens  # wrong: per-seq lengths
// after
cu_q_lens = jnp.concatenate([jnp.zeros(1, jnp.int32), jnp.cumsum(q_lens, dtype=jnp.int32)])

Defensive patterns

Strategy: validation

Validate before calling

cu_q_lens = jnp.concatenate([jnp.zeros(1, jnp.int32), jnp.cumsum(q_lens, dtype=jnp.int32)])
assert cu_q_lens.shape == (page_indices.shape[0] + 1,)

Prevention

Always derive cu_q_lens via concatenate([0], cumsum(q_lens)) instead of hand-building
Check length == num_seqs + 1 before the call

When it happens

Trigger: Passing a cu_q_lens that is too short/long for the batch, or forgetting the terminal total element (passing per-sequence q_lens instead of cumulative sums), when calling ragged_paged_attention.

Common situations: Converting variable-length batched queries into ragged format and forgetting the exclusive-scan final element, or reusing kv_lens as cu_q_lens.

Related errors

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