{"record":{"id":"c07063190faf3a2c","repo":"jax-ml/jax","slug":"not-implemented-num-combined-kv-heads-can-not","errorCode":null,"errorMessage":"Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled.","messagePattern":"Not implemented: (.+?) can not be XLA fully tiled\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":697,"sourceCode":"  bits = dtypes.itemsize_bits(dtype)\n  return 32 // bits\n\n\ndef get_min_heads_per_blk(\n    num_q_heads, num_combined_kv_heads, q_dtype, kv_dtype\n):\n  q_packing = get_dtype_packing(q_dtype)\n  kv_packing = get_dtype_packing(kv_dtype)\n\n  def can_be_xla_fully_tiled(x, packing):\n    if x % packing != 0:\n      return False\n    x //= packing\n    return x in (1, 2, 4, 8) or x % 8 == 0\n\n  # TODO(jevinjiang): support unaligned number of heads!\n  if not can_be_xla_fully_tiled(num_combined_kv_heads, kv_packing):\n    raise ValueError(\n        f\"Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled.\"\n    )\n  assert num_combined_kv_heads % 2 == 0\n  num_kv_heads = num_combined_kv_heads // 2\n  assert num_q_heads % num_kv_heads == 0\n  ratio = num_q_heads // num_kv_heads\n  # TODO(jevinjiang): we can choose smaller tiling for packed type if large\n  # second minor tiling is not on.\n  max_combined_kv_tiling = 8 * kv_packing\n  min_combined_kv_heads = (\n      max_combined_kv_tiling\n      if num_combined_kv_heads % max_combined_kv_tiling == 0\n      else num_combined_kv_heads\n  )\n  min_q_heads = min_combined_kv_heads // 2 * ratio\n  if can_be_xla_fully_tiled(min_q_heads, q_packing):\n    return min_q_heads, min_combined_kv_heads\n  return num_q_heads, num_combined_kv_heads","sourceCodeStart":679,"sourceCodeEnd":715,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L679-L715","documentation":"In get_min_heads_per_blk (used by ragged_paged_attention), the number of combined KV heads (K and V heads packed together, typically 2 * num_kv_heads) must be XLA-tileable: after dividing by the KV packing factor it must be 1, 2, 4, 8, or a multiple of 8. This is a kernel implementation restriction (see the TODO) — non-conforming head counts are not yet supported.","triggerScenarios":"Using KV head counts where 2*num_kv_heads / kv_packing is not in {1,2,4,8} or a multiple of 8, e.g. 6 KV heads with packing 1 (x=12 fails), or 12, 24 combined heads in some packing configurations.","commonSituations":"Running models with unusual GQA head counts (3, 6, 12 KV heads) on this TPU path; changing kv_paging/kernel version where the tiling rule changed.","solutions":["Choose a KV head count such that num_combined_kv_heads // kv_packing is 1, 2, 4, 8, or a multiple of 8 (e.g. 4 or 8 KV heads)","Pad/repeat KV heads to a supported count (jnp.repeat along the head axis) as a workaround","Fall back to a non-Pallas attention implementation until unaligned head counts are supported"],"exampleFix":"// before\nk, v have 6 heads each -> num_combined_kv_heads=12, not tileable\n// after\n# repeat to 8 kv heads\nk = jnp.repeat(k, 2, axis=1); v = jnp.repeat(v, 2, axis=1)  # adjust q head ratio accordingly\n# or use a supported config (num_kv_heads in {1,2,4,8,16,...})","handlingStrategy":"validation","validationCode":"def heads_tileable(num_combined_kv_heads, kv_packing):\n    x = num_combined_kv_heads\n    while x % 2 == 0 and (x // kv_packing) * kv_packing != x:\n        break\n    x = num_combined_kv_heads // kv_packing\n    return x in (1, 2, 4, 8) or x % 8 == 0\nassert heads_tileable(2 * num_kv_heads, kv_packing)","typeGuard":null,"tryCatchPattern":"try:\n    ragged_paged_attention(...)\nexcept ValueError as e:\n    if 'fully tiled' in str(e):\n        attention = fallback_standard_attention  # non-Pallas path\n    else:\n        raise","preventionTips":["Stick to KV head counts of 1,2,4,8,16...","Keep a fallback attention implementation for unsupported head layouts"],"tags":["jax","pallas","tpu","attention","head-tiling","not-implemented"],"backgroundTag":"unsupported-head-count-tpu","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}