{"record":{"id":"68002a5ba4f17230","repo":"jax-ml/jax","slug":"num-q-heads-must-be-divisible-by-num-kv-heads","errorCode":null,"errorMessage":"{num_q_heads=} must be divisible by {num_kv_heads=}","messagePattern":"(.+?) must be divisible by (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":270,"sourceCode":"        \" `max_num_seqs` is `page_indices.shape[0]`.\"\n    )\n  if cu_q_lens.shape != (max_num_seqs + 1,):\n    raise ValueError(\n        f\"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where\"\n        \" `max_num_seqs` is `page_indices.shape[0]`.\"\n    )\n  if (\n      kv_lens.dtype != jnp.int32\n      or page_indices.dtype != jnp.int32\n      or cu_q_lens.dtype != jnp.int32\n  ):\n    raise ValueError(\n        \"The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be\"\n        f\" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},\"\n        f\" {cu_q_lens.dtype=}.\"\n    )\n  if num_q_heads % num_kv_heads != 0:\n    raise ValueError(f\"{num_q_heads=} must be divisible by {num_kv_heads=}\")\n  if sliding_window is not None and sliding_window <= 0:\n    raise ValueError(f\"{sliding_window=} must be positive.\")\n  if soft_cap is not None and soft_cap == 0.0:\n    raise ValueError(f\"{soft_cap=} must not be 0.0.\")\n  if (\n      num_kv_pages_per_block is not None\n      and not 0 < num_kv_pages_per_block <= pages_per_seq\n  ):\n    raise ValueError(\n        f\"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}].\"\n    )\n  if num_queries_per_block is not None and num_queries_per_block <= 0:\n    raise ValueError(f\"{num_queries_per_block=} must be positive.\")\n  if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:\n    raise ValueError(f\"{vmem_limit_bytes=} must be positive.\")\n  del sm_scale  # No constraints on sm_scale.\n  del mask_value  # No consstraints on mask_value.\n","sourceCodeStart":252,"sourceCodeEnd":288,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L252-L288","documentation":"This kernel only implements grouped-query attention (GQA) / multi-query attention where the number of query heads is an integer multiple of the number of KV heads. static_validate_inputs enforces num_q_heads % num_kv_heads == 0 so head-to-head mapping inside the kernel is well defined.","triggerScenarios":"Passing q with a head count not divisible by k/v's head count, e.g. 12 query heads with 8 KV heads.","commonSituations":"Porting a model config to the TPU ragged attention path where head ratios like 8:1 or 7:1 (GQA) are expected; typos in num_heads vs num_kv_heads config; using MHA weights with a partial KV head set.","solutions":["Fix the head configuration so num_q_heads is a multiple of num_kv_heads (e.g. 32 q heads with 8 kv heads)","Repeat/pad KV heads to a divisor of num_q_heads if the model truly has an odd ratio (interleaved repeat like jnp.repeat(kv, ratio, axis=1))","Double-check that q,k,v were not transposed so the head axis is actually axis 1"],"exampleFix":"// before\nq: (seq, 12, d); k/v: (seq, 8, d)  # 12 % 8 != 0\n// after\nk = jnp.repeat(k, 12 // 8 if 12 % 8 == 0 else 1, axis=1)  # better: choose configs like 12 q / 6 kv or 16 q / 8 kv\n# preferred: use num_q_heads=16, num_kv_heads=8","handlingStrategy":"validation","validationCode":"n_q, n_kv = q.shape[1], k.shape[1]\nassert n_q % n_kv == 0, f'{n_q=} not divisible by {n_kv=}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Validate head counts from model config before building the attention call","Standardize on GQA-friendly head ratios (e.g. 32:8, 16:4)"],"tags":["jax","pallas","tpu","gqa","attention","head-dimension"],"backgroundTag":"attention-head-count-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}