jax-ml/jax · error · ValueError

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

Error message

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

What it means

block_kv_compute must itself be divisible by NUM_LANES=8 so the in-kernel softmax running-max tiling works. Even if block_kv is lane-aligned, an unaligned compute block is rejected.

Source

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

    )

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

  q_layout = block_sizes.q_layout

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_kv_compute to a multiple of 8 (64, 128, 256...)
2. Use the library's default block size selection (block_sizes=None)
3. Pin JAX to a version whose recommended block sizes you copied, or regenerate them with the current version's helper

Example fix

// before
BlockSizes(block_kv=256, block_kv_compute=100)
// after
BlockSizes(block_kv=256, block_kv_compute=128)

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

def lane_aligned(x) -> bool: return x is None or x % 8 == 0

When it happens

Trigger: block_sizes=BlockSizes(block_kv=128, block_kv_compute=100) — 128 % 8 == 0 but 100 % 8 != 0.

Common situations: Fine-tuning compute tiling for TPU v5e/v6e where smaller compute blocks improve pipelining; using legacy block sizes from before this constraint existed.

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=} must be a multiple of {bkv_compute=}.
• {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/d12e045497b30408. Report an issue: GitHub.