jax-ml/jax · error · ValueError

{bkv=} must be a multiple of {bkv_compute=}.

Error message

{bkv=} must be a multiple of {bkv_compute=}.

What it means

The KV block size (block_kv) must be divisible by the KV compute block size (block_kv_compute), because the kernel splits each memory block into whole compute blocks. If bkv % bkv_compute != 0 it raises this ValueError.

Source

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

        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:
    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,)}"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make block_kv a multiple of block_kv_compute (simplest: set them equal, e.g. both 128)
2. Or choose block_kv_compute = block_kv // 2 for a valid split
3. Pass block_sizes=None to use library defaults

Example fix

// before
BlockSizes(block_kv=128, block_kv_compute=96)
// after
BlockSizes(block_kv=128, block_kv_compute=128)

assert block_sizes.block_kv % (block_sizes.block_kv_compute or block_sizes.block_kv) == 0

def blocks_divide(bs) -> bool:
    c = bs.block_kv_compute or bs.block_kv
    return bs.block_kv % c == 0

When it happens

Trigger: Passing block_sizes=BlockSizes(block_kv=128, block_kv_compute=96) or any pair where the compute block does not evenly divide the memory block.

Common situations: Hand-tuning block sizes for TPU memory (HBM) vs compute tradeoffs; using configs generated for a different TPU generation.

Related errors

• 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_compute=} must be a multiple of {NUM_LANES}.
• {bq=} should not be greater than {q_seq_len=}

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