jax-ml/jax · error · ValueError

Invalid shape for kv segment_ids: {segment_ids.kv.shape}. Ex

Error message

Invalid shape for kv segment_ids: {segment_ids.kv.shape}. Expected: {(kv_seq_len,)}

What it means

The KV segment ids must be a 1-D array of length kv_seq_len, matching the key/value sequence length exactly. Extra elements (e.g. covering a padded cache) or a batched shape cause this ValueError.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:984

  assert bkv_compute is not None
  if bkv % bkv_compute:
    raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.")
  if bkv_compute % NUM_LANES:
    raise ValueError(f"{bkv_compute=} must be a multiple of {NUM_LANES}.")

  kv_seq_len = k.shape[kv_seq_len_dimension]

  q_heads_per_kv_head = num_q_heads // num_kv_heads

  if segment_ids is not None:
    if segment_ids.q.shape != (q_seq_len,):
      raise ValueError(
          "Invalid shape for q segment_ids: "
          f"{segment_ids.q.shape}. Expected: {(q_seq_len,)}"
      )
    if segment_ids.kv.shape != (kv_seq_len,):
      raise ValueError(
          "Invalid shape for kv segment_ids: "
          f"{segment_ids.kv.shape}. Expected: {(kv_seq_len,)}"
      )

  q_layout = block_sizes.q_layout
  def q_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):
    del j, data_next_ref, mask_next_ref, block_mask_ref
    return from_head_minor((h, i, 0), q_layout)
  def out_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):
    del j, data_next_ref, mask_next_ref, block_mask_ref
    return h, i, 0

  k_layout = block_sizes.k_layout
  def k_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):
    next_j, *_ = _next_nonzero(
        h, i, j, data_next_ref, block_mask_ref, mask_next_ref
    )
    prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Slice kv ids to the current sequence: segment_ids.kv[:kv_seq_len] where kv_seq_len = k.shape[kv_seq_len_dimension]
2. Keep kv ids in sync whenever you append to the cache
3. Check both q and kv segment id shapes before the call

Example fix

// before
seg = SegmentIds(q_ids, kv_ids)  # kv_ids covers padded cache
// after
seg = SegmentIds(q_ids, kv_ids[:k.shape[1]])

assert segment_ids.kv.shape == (kv_seq_len,), f'{segment_ids.kv.shape} != {(kv_seq_len,)}'

def kv_ids_valid(ids, kv_seq_len) -> bool:
    return ids.ndim == 1 and ids.shape[0] == kv_seq_len

When it happens

Trigger: Passing segment_ids.kv of shape [kv_seq_len + padding] when using a preallocated KV cache longer than the actual sequence, or a 2-D batched array.

Common situations: Using paged/KV caches with padding where kv cache length != current sequence length; autoregressive decoding with growing caches.

Related errors

• Invalid shape for q segment_ids: {segment_ids.q.shape}. Expe
• block_kv must be a multiple of {NUM_LANES}
• block_q must be a multiple of {NUM_LANES}
• Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instea
• Expected 'key' head dimension to be: {head_dim_qk}. Instead

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