jax-ml/jax · error · ValueError

Not implemented: {num_combined_kv_heads=} can not be XLA ful

Error message

Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled.

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py:697

  bits = dtypes.itemsize_bits(dtype)
  return 32 // bits


def get_min_heads_per_blk(
    num_q_heads, num_combined_kv_heads, q_dtype, kv_dtype
):
  q_packing = get_dtype_packing(q_dtype)
  kv_packing = get_dtype_packing(kv_dtype)

  def can_be_xla_fully_tiled(x, packing):
    if x % packing != 0:
      return False
    x //= packing
    return x in (1, 2, 4, 8) or x % 8 == 0

  # TODO(jevinjiang): support unaligned number of heads!
  if not can_be_xla_fully_tiled(num_combined_kv_heads, kv_packing):
    raise ValueError(
        f"Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled."
    )
  assert num_combined_kv_heads % 2 == 0
  num_kv_heads = num_combined_kv_heads // 2
  assert num_q_heads % num_kv_heads == 0
  ratio = num_q_heads // num_kv_heads
  # TODO(jevinjiang): we can choose smaller tiling for packed type if large
  # second minor tiling is not on.
  max_combined_kv_tiling = 8 * kv_packing
  min_combined_kv_heads = (
      max_combined_kv_tiling
      if num_combined_kv_heads % max_combined_kv_tiling == 0
      else num_combined_kv_heads
  )
  min_q_heads = min_combined_kv_heads // 2 * ratio
  if can_be_xla_fully_tiled(min_q_heads, q_packing):
    return min_q_heads, min_combined_kv_heads
  return num_q_heads, num_combined_kv_heads

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. 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)
  2. Pad/repeat KV heads to a supported count (jnp.repeat along the head axis) as a workaround
  3. Fall back to a non-Pallas attention implementation until unaligned head counts are supported

Example fix

// before
k, v have 6 heads each -> num_combined_kv_heads=12, not tileable
// after
# repeat to 8 kv heads
k = jnp.repeat(k, 2, axis=1); v = jnp.repeat(v, 2, axis=1)  # adjust q head ratio accordingly
# or use a supported config (num_kv_heads in {1,2,4,8,16,...})
Defensive patterns

Strategy: validation

Validate before calling

def heads_tileable(num_combined_kv_heads, kv_packing):
    x = num_combined_kv_heads
    while x % 2 == 0 and (x // kv_packing) * kv_packing != x:
        break
    x = num_combined_kv_heads // kv_packing
    return x in (1, 2, 4, 8) or x % 8 == 0
assert heads_tileable(2 * num_kv_heads, kv_packing)

Try / catch

try:
    ragged_paged_attention(...)
except ValueError as e:
    if 'fully tiled' in str(e):
        attention = fallback_standard_attention  # non-Pallas path
    else:
        raise

Prevention

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c07063190faf3a2c. Report an issue: GitHub.