{"record":{"id":"e845ade5ac443960","repo":"jax-ml/jax","slug":"bq-should-not-be-greater-than-q-seq-len","errorCode":null,"errorMessage":"{bq=} should not be greater than {q_seq_len=}","messagePattern":"(.+?) should not be greater than (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":1446,"sourceCode":"    mask_value: float,\n    attn_logits_soft_cap: float | None,\n    q_layout: QKVLayout,\n    k_layout: QKVLayout,\n    v_layout: QKVLayout,\n    mask_function: MaskFunctionType | None,\n    interpret: bool,\n):\n  num_q_heads, q_seq_len, head_dim_qk = q.shape\n  head_dim_v = v.shape[-1]\n  if is_mqa:\n    kv_seq_len = k.shape[0]\n    num_kv_heads = 1\n  else:\n    kv_seq_len = k.shape[1]\n    num_kv_heads = k.shape[0]\n\n  if bq > q_seq_len:\n    raise ValueError(\n        f\"{bq=} should not be greater than {q_seq_len=}\")\n  if bkv > kv_seq_len:\n    raise ValueError(\n        f\"{bkv=} should not be greater than {kv_seq_len=}\")\n\n  if not is_mqa and num_q_heads % num_kv_heads != 0:\n    raise ValueError(\n        f\"In MHA, expected number of 'key' heads ({num_kv_heads}) to be a\"\n        f\" multiple of the number of 'query' heads ({num_q_heads})\"\n    )\n\n  if k.shape[:-1] != v.shape[:-1]:\n    raise ValueError(\n        f\"Expected 'key' {k.shape} and 'value' {v.shape} to have the same \"\n        \"leading dimensions.\"\n    )\n\n  if bkv % NUM_LANES:","sourceCodeStart":1428,"sourceCodeEnd":1464,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L1428-L1464","documentation":"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.","triggerScenarios":"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.","commonSituations":"Autoregressive decoding where the query length is 1–8 tokens while block sizes were tuned for prefill (512+); unit tests with tiny sequence lengths.","solutions":["Reduce block_q to <= q_seq_len (e.g. block_q=min(128, q_seq_len))","Pad the query sequence to at least block_q","For decode workloads use a decode-specialized kernel (paged attention) instead of splash attention"],"exampleFix":"// before\nblock_sizes=BlockSizes(block_q=512, ...)\nloss = grad_fn(q_len_256)\n// after\nblock_sizes=BlockSizes(block_q=min(256, q_seq_len), ...)\nloss = grad_fn(q_len_256)","handlingStrategy":"validation","validationCode":"block_sizes = BlockSizes(block_q=min(block_q, q_seq_len), block_kv=block_kv, ...)","typeGuard":"def bq_fits(bq: int, q_seq_len: int) -> bool: return bq <= q_seq_len","tryCatchPattern":null,"preventionTips":["Compute block sizes dynamically from actual sequence lengths","Use decode kernels for single-token queries"],"tags":["jax","pallas","tpu","splash-attention","block-size","sequence-length"],"backgroundTag":"block-size-constraint-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}