{"record":{"id":"3b6e044b3b0f0326","repo":"jax-ml/jax","slug":"invalid-shape-for-kv-segment-ids-segment-ids-kv","errorCode":null,"errorMessage":"Invalid shape for kv segment_ids: {segment_ids.kv.shape}. Expected: {(kv_seq_len,)}","messagePattern":"Invalid shape for kv segment_ids: (.+?)\\. Expected: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":984,"sourceCode":"\n  assert bkv_compute is not None\n  if bkv % bkv_compute:\n    raise ValueError(f\"{bkv=} must be a multiple of {bkv_compute=}.\")\n  if bkv_compute % NUM_LANES:\n    raise ValueError(f\"{bkv_compute=} must be a multiple of {NUM_LANES}.\")\n\n  kv_seq_len = k.shape[kv_seq_len_dimension]\n\n  q_heads_per_kv_head = num_q_heads // num_kv_heads\n\n  if segment_ids is not None:\n    if segment_ids.q.shape != (q_seq_len,):\n      raise ValueError(\n          \"Invalid shape for q segment_ids: \"\n          f\"{segment_ids.q.shape}. Expected: {(q_seq_len,)}\"\n      )\n    if segment_ids.kv.shape != (kv_seq_len,):\n      raise ValueError(\n          \"Invalid shape for kv segment_ids: \"\n          f\"{segment_ids.kv.shape}. Expected: {(kv_seq_len,)}\"\n      )\n\n  q_layout = block_sizes.q_layout\n  def q_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):\n    del j, data_next_ref, mask_next_ref, block_mask_ref\n    return from_head_minor((h, i, 0), q_layout)\n  def out_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):\n    del j, data_next_ref, mask_next_ref, block_mask_ref\n    return h, i, 0\n\n  k_layout = block_sizes.k_layout\n  def k_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None):\n    next_j, *_ = _next_nonzero(\n        h, i, j, data_next_ref, block_mask_ref, mask_next_ref\n    )\n    prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),)","sourceCodeStart":966,"sourceCodeEnd":1002,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L966-L1002","documentation":"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.","triggerScenarios":"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.","commonSituations":"Using paged/KV caches with padding where kv cache length != current sequence length; autoregressive decoding with growing caches.","solutions":["Slice kv ids to the current sequence: segment_ids.kv[:kv_seq_len] where kv_seq_len = k.shape[kv_seq_len_dimension]","Keep kv ids in sync whenever you append to the cache","Check both q and kv segment id shapes before the call"],"exampleFix":"// before\nseg = SegmentIds(q_ids, kv_ids)  # kv_ids covers padded cache\n// after\nseg = SegmentIds(q_ids, kv_ids[:k.shape[1]])","handlingStrategy":"validation","validationCode":"assert segment_ids.kv.shape == (kv_seq_len,), f'{segment_ids.kv.shape} != {(kv_seq_len,)}'","typeGuard":"def kv_ids_valid(ids, kv_seq_len) -> bool:\n    return ids.ndim == 1 and ids.shape[0] == kv_seq_len","tryCatchPattern":null,"preventionTips":["Slice kv ids to the live sequence length when using padded caches","Keep ids and cache lengths updated together"],"tags":["jax","pallas","tpu","splash-attention","segment-ids","shape"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}