jax-ml/jax · error · ValueError

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

Error message

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

What it means

In the backward dq kernel, block_kv must be divisible by NUM_LANES=8 to keep lane-aligned softmax reductions correct during the gradient computation. The source notes block_kv_compute alignment is a TODO, but bkv itself is enforced.

Source

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

        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
  # multiple of NUM_LANES.

  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[-1]
  else:
    grid_width = kv_seq_len // bkv

  grid = (num_q_heads, q_seq_len // bq, grid_width)

  def o_index_map(h, i, *_):
    return h, i, 0

  o_spec = pl.BlockSpec((None, bq, head_dim_v), o_index_map)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_kv to a multiple of 8 (64, 128, 256)
2. Use default block sizes (block_sizes=None) which are lane-aligned
3. Validate all block sizes against NUM_LANES=8 in your config loader

Example fix

// before
BlockSizes(block_kv=112)
// after
BlockSizes(block_kv=128)

assert block_sizes.block_kv % 8 == 0

def lane_aligned(x: int) -> bool: return x % 8 == 0

When it happens

Trigger: jax.grad through splash attention with block_sizes.block_kv=100 or any value not divisible by 8 (e.g. tuned for HBM size 7*16=112).

Common situations: Hand-tuned block sizes that satisfy the forward constraints but forgot the backward's NUM_LANES constraint on block_kv.

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=}.
• {bkv_compute=} must be a multiple of {NUM_LANES}.

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