jax-ml/jax · error · ValueError

Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instea

Error message

Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a {len(k.shape)}-dim one.

What it means

Splash Attention expects the 'key' tensor to have exactly expected_kv_rank dimensions (3 for MHA: [num_kv_heads, kv_seq_len, head_dim]; 2 for MQA: [kv_seq_len, head_dim]). A different rank means the layout assumption of the kernel is broken, so it fails fast.

Source

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

    num_kv_heads = 1
  else:
    expected_kv_rank = 3
    kv_head_dimension = 2
    kv_seq_len_dimension = 1
    num_kv_heads = k.shape[0]

  partial_mask_blocks = fwd_mask_info.partial_mask_blocks
  if (
      partial_mask_blocks is not None
      and jnp.dtype(partial_mask_blocks.dtype) != np.bool_
  ):
    raise ValueError(
        "partial_mask_blocks must be of type np.bool_ but got"
        f" {partial_mask_blocks.dtype}"
    )

  if len(k.shape) != expected_kv_rank:
    raise ValueError(
        f"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a"
        f" {len(k.shape)}-dim one."
    )

  if k.shape[kv_head_dimension] != head_dim_qk:
    raise ValueError(
        f"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:"
        f" {k.shape[kv_head_dimension]}."
    )

  if not is_mqa and num_q_heads % num_kv_heads != 0:
    raise ValueError(
        f"In MHA, expected number of 'key' heads ({num_kv_heads}) to be a"
        f" multiple of the number of 'query' heads ({num_q_heads})"
    )

  if k.shape[:-1] != v.shape[:-1]:
    raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the batch dimension: splash attention is single-sequence per call, pass k[num_kv_heads, kv_seq_len, head_dim] (or [kv_seq_len, head_dim] for MQA)
2. Check is_mqa inference: a 2-D k is treated as MQA; reshape accordingly
3. If you have a batch, vmap the attention function over the batch axis instead of passing it inside k

Example fix

// before
out, _ = attn(q[0], k[0], v[0])  # k[0] still [batch, heads, seq, dim] misuse
# or k = k.reshape(batch, kv_heads, seq, dim)
// after
attn_fn = jax.vmap(splash_fn)
out = attn_fn(q, k, v)  # k: [batch, kv_heads, seq, dim], vmapped

assert k.ndim == (2 if is_mqa else 3), f'key must be {2 if is_mqa else 3}-D, got {k.ndim}'

def key_rank_ok(k, is_mqa: bool) -> bool:
    return k.ndim == (2 if is_mqa else 3)

When it happens

Trigger: Passing k with shape [batch, num_kv_heads, kv_seq_len, head_dim] (4-D, batched like in flash attention) to the unbatched splash API, or passing an MQA-shaped 2-D tensor while declaring MHA, or vice versa.

Common situations: Porting code from jax.nn.dot_product_attention or GPU flash attention where a batch dim is expected; loading K from a checkpoint with an extra leading axis.

Related errors

• Expected 'key' head dimension to be: {head_dim_qk}. Instead
• In MHA, expected number of 'key' heads ({num_kv_heads}) to b
• Expected 'key' {k.shape} and 'value' {v.shape} to have the s
• Invalid shape for q segment_ids: {segment_ids.q.shape}. Expe
• Invalid shape for kv segment_ids: {segment_ids.kv.shape}. Ex

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