{"record":{"id":"a46e395c6e8da3f6","repo":"sgl-project/sglang","slug":"k-shape-must-be-num-tokens-num-kv-heads-head-di","errorCode":null,"errorMessage":"k shape must be [num_tokens, num_kv_heads, head_dim], got {k.shape}","messagePattern":"k shape must be \\[num_tokens, num_kv_heads, head_dim\\], got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/aot/python/sgl_kernel/metal.py","lineNumber":82,"sourceCode":"        raise ValueError(\"rope_pool_fused expects q/k/v to be 3-D\")\n    if positions.ndim != 1 or slots.ndim != 1:\n        raise ValueError(\"rope_pool_fused expects positions/slots to be 1-D\")\n    if k_pool.ndim != 3 or v_pool.ndim != 3:\n        raise ValueError(\"rope_pool_fused expects pool tensors to be 3-D\")\n    q_shape = tuple(q.shape)\n    k_shape = tuple(k.shape)\n    v_shape = tuple(v.shape)\n    positions_shape = tuple(positions.shape)\n    slots_shape = tuple(slots.shape)\n    k_pool_shape = tuple(k_pool.shape)\n    v_pool_shape = tuple(v_pool.shape)\n\n    if q_shape != (q_shape[0], num_qo_heads, head_dim):\n        raise ValueError(\n            \"q shape must be [num_tokens, num_qo_heads, head_dim], \" f\"got {q.shape}\"\n        )\n    if k_shape != (q_shape[0], num_kv_heads, head_dim):\n        raise ValueError(\n            \"k shape must be [num_tokens, num_kv_heads, head_dim], \" f\"got {k.shape}\"\n        )\n    if v_shape != k_shape:\n        raise ValueError(f\"v shape must match k shape, got {v.shape} vs {k.shape}\")\n    if positions_shape != (q_shape[0],) or slots_shape != (q_shape[0],):\n        raise ValueError(\"positions/slots must have one entry per token\")\n    if k_pool_shape[1:] != (num_kv_heads, head_dim):\n        raise ValueError(f\"k_pool has incompatible shape {k_pool.shape}\")\n    if v_pool_shape != k_pool_shape:\n        raise ValueError(\n            f\"v_pool shape must match k_pool shape, got {v_pool.shape} vs {k_pool.shape}\"\n        )\n    if q.dtype != k.dtype or q.dtype != v.dtype:\n        raise ValueError(\"q/k/v dtypes must match\")\n    if k_pool.dtype != q.dtype or v_pool.dtype != q.dtype:\n        raise ValueError(\"pool dtypes must match q/k/v dtype\")\n\n    return _metal.rope_pool_fused(","sourceCodeStart":64,"sourceCodeEnd":100,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/aot/python/sgl_kernel/metal.py#L64-L100","documentation":"rope_pool_fused validates that k has shape [num_tokens, num_kv_heads, head_dim], sharing num_tokens with q and using the declared KV head count. This error fires when k's shape disagrees — typically wrong number of KV heads or head_dim, or a different token count than q.","triggerScenarios":"Passing k with shape[1] != num_kv_heads or shape[2] != head_dim, or k.shape[0] != q.shape[0] (e.g. k trimmed/extended differently from q in a prefill path).","commonSituations":"GQA misconfiguration where num_kv_heads doesn't match the K projection width; slicing q and k with different token ranges; using MHA-style k (all heads) with GQA num_kv_heads from config.","solutions":["Ensure k.shape == (q.shape[0], num_kv_heads, head_dim)","Derive num_kv_heads from the model config: num_kv_heads = k.shape[1] and pass that","Align token counts: k = k[: q.shape[0]] if q was sliced"],"exampleFix":"# before\nnum_kv_heads = config.num_attention_heads  # wrong for GQA\n\n# after\nnum_kv_heads = config.num_key_value_heads\nassert k.shape == (q.shape[0], num_kv_heads, head_dim)","handlingStrategy":"validation","validationCode":"assert k.shape == (q.shape[0], num_kv_heads, head_dim), (k.shape, num_kv_heads, head_dim)","typeGuard":"def k_shape_ok(q, k, num_kv_heads, head_dim):\n    return k.shape == (q.shape[0], num_kv_heads, head_dim)","tryCatchPattern":null,"preventionTips":["Use config.num_key_value_heads for GQA models","Slice q and k with the same token range"],"tags":["shape-validation","gqa","rope","metal","sgl-kernel"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}