jax-ml/jax · error · ValueError

{bq=} should not be greater than {q_seq_len=}

Error message

{bq=} should not be greater than {q_seq_len=}

What it means

In the backward dq kernel, the query block size bq must not exceed the actual query sequence length q_seq_len; otherwise grid blocks would read past the end of q. The kernel validates this before tracing.

Source

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

    mask_value: float,
    attn_logits_soft_cap: float | None,
    q_layout: QKVLayout,
    k_layout: QKVLayout,
    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:
    kv_seq_len = k.shape[0]
    num_kv_heads = 1
  else:
    kv_seq_len = k.shape[1]
    num_kv_heads = k.shape[0]

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce block_q to <= q_seq_len (e.g. block_q=min(128, q_seq_len))
2. Pad the query sequence to at least block_q
3. For decode workloads use a decode-specialized kernel (paged attention) instead of splash attention

Example fix

// before
block_sizes=BlockSizes(block_q=512, ...)
loss = grad_fn(q_len_256)
// after
block_sizes=BlockSizes(block_q=min(256, q_seq_len), ...)
loss = grad_fn(q_len_256)

block_sizes = BlockSizes(block_q=min(block_q, q_seq_len), block_kv=block_kv, ...)

def bq_fits(bq: int, q_seq_len: int) -> bool: return bq <= q_seq_len

When it happens

Trigger: Calling the backward pass (jax.grad of splash attention) with block_sizes.block_q=512 but q_seq_len=256; also short sequences during decoding (q_seq_len=1) with default large blocks.

Common situations: Autoregressive decoding where the query length is 1–8 tokens while block sizes were tuned for prefill (512+); unit tests with tiny sequence lengths.

Related errors

• {bkv=} should not be greater than {kv_seq_len=}
• 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=}.

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