{"record":{"id":"9c7ea117cdb23e8c","repo":"xai-org/x-algorithm","slug":"num-q-heads-must-be-divisible-by-and-num-kv-he","errorCode":null,"errorMessage":"{num_q_heads=} must be divisible by and {num_kv_heads=}","messagePattern":"(.+?) must be divisible by and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":103,"sourceCode":"    sm_scale: float = 1.0,\n    cap: float = -1.0,\n    cap_method: str = \"tanh\",\n):\n    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:\n        raise ValueError(f\"q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}\")\n    batch_size, q_seq_len, num_q_heads, head_dim = q.shape\n    _, kv_seq_len, num_kv_heads, _ = k.shape\n    kv_shape = (batch_size, kv_seq_len, num_kv_heads, head_dim)\n    if k.shape != kv_shape:\n        raise ValueError(f\"Expected {k.shape=} to be {kv_shape} (inferred from q)\")\n    if v.shape != kv_shape:\n        raise ValueError(f\"Expected {v.shape=} to be {kv_shape} (inferred from q)\")\n    if (dtype := q.dtype) != k.dtype or dtype != v.dtype:\n        raise ValueError(\n            f\"q, k, and v should all have the same dtype, got: {q.dtype}, {k.dtype}, {v.dtype}\"\n        )\n    if num_q_heads % num_kv_heads:\n        raise ValueError(f\"{num_q_heads=} must be divisible by and {num_kv_heads=}\")\n    q_heads_per_kv_head = num_q_heads // num_kv_heads\n    if head_dim % 64:\n        raise ValueError(f\"{head_dim=} must be divisible by 64\")\n    if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):\n        raise NotImplementedError(f\"Only f16 and bf16 are supported, got dtype: {dtype}\")\n\n    max_concurrent_steps = min(config.max_concurrent_steps, kv_seq_len // config.block_kv)\n    block_q, block_kv = config.block_q, config.block_kv\n    if kv_seq_len % block_kv:\n        raise ValueError(f\"{kv_seq_len=} must be a multiple of {block_kv=}\")\n\n    def kernel(q_ref, k_ref, v_ref, bound_ref, out_ref, lse_ref, scoped):\n        batch = lax.axis_index(\"batch\")\n        q_head = lax.axis_index(\"heads\")\n        q_seq = lax.axis_index(\"q_seq\")\n        smem_buffers, buffer_barriers, consumed_barriers, schedule_barrier = scoped\n        wg_idx = lax.axis_index(\"wg\")\n        qo_smem2, k_smem, v_smem, lse_smem2 = smem_buffers","sourceCodeStart":85,"sourceCodeEnd":121,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L85-L121","documentation":"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.","triggerScenarios":"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).","commonSituations":"Head pruning or distillation that changes q heads only; configs hand-edited to an odd kv-head count; local attention with different head splits.","solutions":["Set num_kv_heads to a divisor of num_q_heads (1 for MQA, num_q_heads for MHA, or an even group size)","Align both projections' head counts after any pruning"],"exampleFix":"# before\nnum_q_heads, num_kv_heads = 12, 8\n# after\nnum_q_heads, num_kv_heads = 12, 4","handlingStrategy":"validation","validationCode":"assert num_q_heads % num_kv_heads == 0, \"GQA requires num_kv_heads | num_q_heads\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive num_kv_heads as num_q_heads // group_size in config","Add a model-config test for head divisibility"],"tags":["jax","attention","gqa","head-dims"],"backgroundTag":"invalid-head-count-config","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}