jax-ml/jax · error · ValueError

Invalid shape for q segment_ids: {segment_ids.q.shape}. Expe

Error message

Invalid shape for q segment_ids: {segment_ids.q.shape}. Expected: {(q_seq_len,)}

What it means

When segment_ids are provided, the query segment ids must be a 1-D array of length q_seq_len (per-sequence, no batch or head axes). Any other shape is rejected before launching the kernel.

Source

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

  if k.shape[:-1] != v.shape[:-1]:
    raise ValueError(
        f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same "
        "leading dimensions."
    )

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a flat 1-D array: segment_ids.q = q_ids[q_seq_len]
2. If batched, vmap over segment ids too (in_axes=(..., 0, 0, SegmentIds(0, 0))) or index per batch manually
3. Verify kv segment ids shape is (kv_seq_len,) at the same time

Example fix

// before
seg = SegmentIds(q=q_ids, kv=kv_ids)  # q_ids: [B, S]
out = fn(q[0], k[0], v[0], segment_ids=seg)
// after
out = fn(q[0], k[0], v[0], segment_ids=SegmentIds(q_ids[0], kv_ids[0]))

assert segment_ids.q.shape == (q_seq_len,), f'{segment_ids.q.shape} != {(q_seq_len,)}'

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

When it happens

Trigger: Passing segment_ids.q with shape [batch, q_seq_len] or [num_heads, q_seq_len] instead of [q_seq_len]; forgetting to index out the batch when using vmap with in_axes configured wrongly.

Common situations: Porting from GPU pipeshard/flash attention where segment ids carry a batch dim; vmap'ing attention but leaving segment ids un-vmapped.

Related errors

• Invalid shape for kv segment_ids: {segment_ids.kv.shape}. Ex
• 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/da6ef01aeef60555. Report an issue: GitHub.