{"record":{"id":"080e91351b9011b2","repo":"jax-ml/jax","slug":"bkv-must-be-a-multiple-of-bkv-compute","errorCode":null,"errorMessage":"{bkv=} must be a multiple of {bkv_compute=}.","messagePattern":"(.+?) must 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":969,"sourceCode":"        f\"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:\"\n        f\" {k.shape[kv_head_dimension]}.\"\n    )\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  assert bkv_compute is not None\n  if bkv % bkv_compute:\n    raise ValueError(f\"{bkv=} must be a multiple of {bkv_compute=}.\")\n  if bkv_compute % NUM_LANES:\n    raise ValueError(f\"{bkv_compute=} must be a multiple of {NUM_LANES}.\")\n\n  kv_seq_len = k.shape[kv_seq_len_dimension]\n\n  q_heads_per_kv_head = num_q_heads // num_kv_heads\n\n  if segment_ids is not None:\n    if segment_ids.q.shape != (q_seq_len,):\n      raise ValueError(\n          \"Invalid shape for q segment_ids: \"\n          f\"{segment_ids.q.shape}. Expected: {(q_seq_len,)}\"\n      )\n    if segment_ids.kv.shape != (kv_seq_len,):\n      raise ValueError(\n          \"Invalid shape for kv segment_ids: \"\n          f\"{segment_ids.kv.shape}. Expected: {(kv_seq_len,)}\"\n      )","sourceCodeStart":951,"sourceCodeEnd":987,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L951-L987","documentation":"The KV block size (block_kv) must be divisible by the KV compute block size (block_kv_compute), because the kernel splits each memory block into whole compute blocks. If bkv % bkv_compute != 0 it raises this ValueError.","triggerScenarios":"Passing block_sizes=BlockSizes(block_kv=128, block_kv_compute=96) or any pair where the compute block does not evenly divide the memory block.","commonSituations":"Hand-tuning block sizes for TPU memory (HBM) vs compute tradeoffs; using configs generated for a different TPU generation.","solutions":["Make block_kv a multiple of block_kv_compute (simplest: set them equal, e.g. both 128)","Or choose block_kv_compute = block_kv // 2 for a valid split","Pass block_sizes=None to use library defaults"],"exampleFix":"// before\nBlockSizes(block_kv=128, block_kv_compute=96)\n// after\nBlockSizes(block_kv=128, block_kv_compute=128)","handlingStrategy":"validation","validationCode":"assert block_sizes.block_kv % (block_sizes.block_kv_compute or block_sizes.block_kv) == 0","typeGuard":"def blocks_divide(bs) -> bool:\n    c = bs.block_kv_compute or bs.block_kv\n    return bs.block_kv % c == 0","tryCatchPattern":null,"preventionTips":["Set block_kv_compute equal to block_kv unless profiling says otherwise","Validate BlockSizes in a factory function"],"tags":["jax","pallas","tpu","splash-attention","block-size"],"backgroundTag":"block-size-constraint-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}