jax-ml/jax · error · ValueError

{bkv=} should not be greater than {kv_seq_len=}

Error message

{bkv=} should not be greater than {kv_seq_len=}

What it means

In the backward dq kernel, the KV block size bkv must not exceed the KV sequence length kv_seq_len, since blocks would index out of bounds of k/v. Checked alongside the analogous q constraint.

Source

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

    k_layout: QKVLayout,
    v_layout: QKVLayout,
    mask_function: MaskFunctionType | None,
    interpret: bool,
):
  num_q_heads, q_seq_len, head_dim_qk = q.shape
  head_dim_v = v.shape[-1]
  if is_mqa:
    kv_seq_len = k.shape[0]
    num_kv_heads = 1
  else:
    kv_seq_len = k.shape[1]
    num_kv_heads = k.shape[0]

  if bq > q_seq_len:
    raise ValueError(
        f"{bq=} should not be greater than {q_seq_len=}")
  if bkv > kv_seq_len:
    raise ValueError(
        f"{bkv=} should not be greater than {kv_seq_len=}")

  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."
    )

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

  # TODO(amagni/sharadmv): when adding block_compute, make sure that is a

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_kv <= kv_seq_len (e.g. min(block_kv, kv_seq_len))
2. Pad K/V to a multiple of block_kv if padding is acceptable
3. Use decode-optimized kernels for short sequences

Example fix

// before
BlockSizes(block_q=128, block_kv=1024)  # kv_seq_len=512
// after
BlockSizes(block_q=128, block_kv=512)

block_sizes = BlockSizes(block_kv=min(block_kv, kv_seq_len), ...)

def bkv_fits(bkv: int, kv_seq_len: int) -> bool: return bkv <= kv_seq_len

When it happens

Trigger: Backward pass with block_kv=1024 while kv_seq_len=512; short KV caches during decode; tiny test sequences.

Common situations: Reusing prefill block sizes during decode with a small KV cache; small unit tests with seq_len < block sizes.

Related errors

• {bq=} should not be greater than {q_seq_len=}
• block_kv must be a multiple of {NUM_LANES}
• block_q must be a multiple of {NUM_LANES}
• {bkv_compute=} should be a multiple of {NUM_LANES}
• {bkv=} must be a multiple of {bkv_compute=}.

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