jax-ml/jax · error · ValueError

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

Error message

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

What it means

Splash Attention backward validates that the compute tile for KV (bkv_compute) does not exceed the memory tile bkv. The kernel splits the KV block into compute sub-blocks, so the sub-block must fit inside the block.

Source

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

    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:
    num_kv_heads, kv_seq_len = 1, k.shape[0]
  else:
    num_kv_heads, kv_seq_len, _ = k.shape

  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 bkv_compute > bkv:
    raise ValueError(
        f"{bkv_compute=} should not be greater than {bkv=}")
  if bkv % bkv_compute:
    raise ValueError(
        f"{bkv=} should be a multiple of {bkv_compute=}")

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

  q_heads_per_kv_head = num_q_heads // num_kv_heads

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_kv_dkv_compute <= block_kv_dkv
2. Prefer leaving bkv_compute at its default (often None) so it is derived automatically
3. Validate the full BlockSizes tuple together rather than overriding single fields

Example fix

// before
BlockSizes(block_kv_dkv=128, block_kv_dkv_compute=256)
// after
BlockSizes(block_kv_dkv=256, block_kv_dkv_compute=128)

assert bs.block_kv_dkv_compute is None or bs.block_kv_dkv_compute <= bs.block_kv_dkv

When it happens

Trigger: Setting BlockSizes.block_kv_dkv_compute > BlockSizes.block_kv_dkv when configuring backward blocks for splash attention on TPU.

Common situations: Hand-tuning block sizes for TPU memory; copying block-size configs from a different kernel version where defaults differ; partial overrides that leave an inconsistent combination.

Related errors

• {bkv=} should be a multiple of {bkv_compute=}
• Invalid memory space: {memory_space!r}
• {num_queries_per_block=} must be positive.
• {vmem_limit_bytes=} must be positive.
• Higher-order AD not supported.

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