xai-org/x-algorithm · error · ValueError

{num_q_heads=} must be divisible by and {num_kv_heads=}

Error message

{num_q_heads=} must be divisible by and {num_kv_heads=}

What it means

For grouped-query attention the number of query heads must be divisible by the number of kv heads so each kv head serves an integer number of query heads (q_heads_per_kv_head = num_q_heads // num_kv_heads). The non-divisible remainder makes this invalid GQA/MQA.

Source

Thrown at phoenix/xrex/pallas/ranker_attention_fa3.py:103

    sm_scale: float = 1.0,
    cap: float = -1.0,
    cap_method: str = "tanh",
):
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError(f"q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}")
    batch_size, q_seq_len, num_q_heads, head_dim = q.shape
    _, kv_seq_len, num_kv_heads, _ = k.shape
    kv_shape = (batch_size, kv_seq_len, num_kv_heads, head_dim)
    if k.shape != kv_shape:
        raise ValueError(f"Expected {k.shape=} to be {kv_shape} (inferred from q)")
    if v.shape != kv_shape:
        raise ValueError(f"Expected {v.shape=} to be {kv_shape} (inferred from q)")
    if (dtype := q.dtype) != k.dtype or dtype != v.dtype:
        raise ValueError(
            f"q, k, and v should all have the same dtype, got: {q.dtype}, {k.dtype}, {v.dtype}"
        )
    if num_q_heads % num_kv_heads:
        raise ValueError(f"{num_q_heads=} must be divisible by and {num_kv_heads=}")
    q_heads_per_kv_head = num_q_heads // num_kv_heads
    if head_dim % 64:
        raise ValueError(f"{head_dim=} must be divisible by 64")
    if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):
        raise NotImplementedError(f"Only f16 and bf16 are supported, got dtype: {dtype}")

    max_concurrent_steps = min(config.max_concurrent_steps, kv_seq_len // config.block_kv)
    block_q, block_kv = config.block_q, config.block_kv
    if kv_seq_len % block_kv:
        raise ValueError(f"{kv_seq_len=} must be a multiple of {block_kv=}")

    def kernel(q_ref, k_ref, v_ref, bound_ref, out_ref, lse_ref, scoped):
        batch = lax.axis_index("batch")
        q_head = lax.axis_index("heads")
        q_seq = lax.axis_index("q_seq")
        smem_buffers, buffer_barriers, consumed_barriers, schedule_barrier = scoped
        wg_idx = lax.axis_index("wg")
        qo_smem2, k_smem, v_smem, lse_smem2 = smem_buffers

View on GitHub (pinned to 24c60942c5)

Solutions

1. Set num_kv_heads to a divisor of num_q_heads (1 for MQA, num_q_heads for MHA, or an even group size)
2. Align both projections' head counts after any pruning

Example fix

# before
num_q_heads, num_kv_heads = 12, 8
# after
num_q_heads, num_kv_heads = 12, 4

assert num_q_heads % num_kv_heads == 0, "GQA requires num_kv_heads | num_q_heads"

When it happens

Trigger: num_q_heads=12 with num_kv_heads=8, or any pair where num_q_heads % num_kv_heads != 0 (e.g. after head-pruning of q but not kv).

Common situations: Head pruning or distillation that changes q heads only; configs hand-edited to an odd kv-head count; local attention with different head splits.

Related errors

• Expected {k.shape=} to be {kv_shape} (inferred from q)
• Please override this method for specific attention impl.
• Mask dimensionality {mask.ndim} must match logits dimensiona
• Invalid backward pass implementation: {backward_pass_impl}
• cap_method must be in [tanh, soft_sign, none], got {cap_meth

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/9c7ea117cdb23e8c. Report an issue: GitHub.