jax-ml/jax · error · ValueError

Expected 'key' head dimension to be: {head_dim_qk}. Instead

Error message

Expected 'key' head dimension to be: {head_dim_qk}. Instead got: {k.shape[kv_head_dimension]}.

What it means

The head dimension of the 'key' tensor (k.shape[kv_head_dimension]) must equal the query head dimension head_dim_qk that the attention function was built for. Mismatched QK head dims make the dot product undefined, so the kernel validates it.

Source

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

  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(
        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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Project K (and V) to the same head dimension as Q before calling splash attention, e.g. with an einsum/Linear
Fix the head_dim argument used when constructing the attention function to match your tensors
Check that you sliced k = fused_qkv[..., 2*hd:3*hd] correctly if using a fused projection

Example fix

// before
k = fused[..., q_hd:q_hd+kv_hd]  # kv_hd != q_hd
// after
k_proj = hk.Linear(q_hd)(k)
attn(q, k_proj, v_proj)

Defensive patterns

Strategy: validation

Validate before calling

assert k.shape[-1] == head_dim_qk == q.shape[-1], 'QK head dims must match'

Type guard

def head_dims_match(q, k, head_dim) -> bool:
    return q.shape[-1] == k.shape[-1] == head_dim

Prevention

Project K/V to the Q head dim for GQA models with distinct KV dims
Centralize head-dim constants in config

When it happens

Trigger: Building splash attention with head_dim=128 but passing k whose last dim is 64 (e.g. from a GQA checkpoint with a different KV head size than Q head size beyond supported remapping).

Common situations: Using architectures where KV heads are quantized/projected to a different dim (e.g. MLA-style 512-dim KV with 128-dim Q); passing the wrong slice of a fused qkv projection.

Related errors

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