{"record":{"id":"ee699540ebbfdc23","repo":"jax-ml/jax","slug":"bkv-should-be-a-multiple-of-bkv-compute","errorCode":null,"errorMessage":"{bkv=} should be a multiple of {bkv_compute=}","messagePattern":"(.+?) should be a multiple of (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":1911,"sourceCode":"):\n  num_q_heads, q_seq_len, head_dim_qk = q.shape\n  head_dim_v = v.shape[-1]\n  if is_mqa:\n    num_kv_heads, kv_seq_len = 1, k.shape[0]\n  else:\n    num_kv_heads, kv_seq_len, _ = k.shape\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  if bkv_compute > bkv:\n    raise ValueError(\n        f\"{bkv_compute=} should not be greater than {bkv=}\")\n  if bkv % bkv_compute:\n    raise ValueError(\n        f\"{bkv=} should be a multiple of {bkv_compute=}\")\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  q_heads_per_kv_head = num_q_heads // num_kv_heads\n\n  if mask_info.data_next is not None:\n    grid_width = mask_info.data_next.shape[-2]","sourceCodeStart":1893,"sourceCodeEnd":1929,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L1893-L1929","documentation":"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.","triggerScenarios":"Configuring BlockSizes for splash attention backward where block_kv_dkv % block_kv_dkv_compute != 0, e.g. bkv=192 with bkv_compute=128.","commonSituations":"Tuning block sizes for TPU HBM/VMEM tradeoffs with non-power-of-two values; mixing configurations from different kernel revisions.","solutions":["Choose bkv_compute that divides bkv evenly (e.g. 128/128, 256/128, 256/64)","Leave bkv_compute as None to use the default derived from bkv","Double-check that any auto-derived default was not disabled"],"exampleFix":"// before\nBlockSizes(block_kv_dkv=192, block_kv_dkv_compute=128)\n// after\nBlockSizes(block_kv_dkv=256, block_kv_dkv_compute=128)","handlingStrategy":"validation","validationCode":"assert bs.block_kv_dkv_compute is None or bs.block_kv_dkv % bs.block_kv_dkv_compute == 0","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use power-of-two block sizes","Validate all block invariants in one check_block_sizes() helper"],"tags":["jax","tpu","pallas","splash-attention","block-sizes","divisibility"],"backgroundTag":"block-size-divisibility","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}