jax-ml/jax · error · ValueError

{bkv=} should be a multiple of {bkv_compute=}

Error message

{bkv=} should be a multiple of {bkv_compute=}

What it means

The KV memory block (bkv) must be an integer multiple of the KV compute block (bkv_compute), because the kernel iterates compute sub-tiles to cover the memory tile exactly.

Source

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

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

  if mask_info.data_next is not None:
    grid_width = mask_info.data_next.shape[-2]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Choose bkv_compute that divides bkv evenly (e.g. 128/128, 256/128, 256/64)
2. Leave bkv_compute as None to use the default derived from bkv
3. Double-check that any auto-derived default was not disabled

Example fix

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

assert bs.block_kv_dkv_compute is None or bs.block_kv_dkv % bs.block_kv_dkv_compute == 0

When it happens

Trigger: Configuring BlockSizes for splash attention backward where block_kv_dkv % block_kv_dkv_compute != 0, e.g. bkv=192 with bkv_compute=128.

Common situations: Tuning block sizes for TPU HBM/VMEM tradeoffs with non-power-of-two values; mixing configurations from different kernel revisions.

Related errors

• {bkv_compute=} should not be greater than {bkv=}
• Higher-order AD not supported.
• Block sizes for dq kernel are not needed with a fused kernel
• Mask function must return a boolean-valued array, but got: {
• block_kv must be a multiple of {NUM_LANES}

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