{"record":{"id":"088cdb14f59c6b97","repo":"sgl-project/sglang","slug":"rope-pool-fused-expects-q-k-v-to-be-3-d","errorCode":null,"errorMessage":"rope_pool_fused expects q/k/v to be 3-D","messagePattern":"rope_pool_fused expects q/k/v to be 3-D","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/aot/python/sgl_kernel/metal.py","lineNumber":64,"sourceCode":"    rope_base: float,\n) -> tuple[mx.array, mx.array, mx.array, mx.array]:\n    \"\"\"Apply NeoX RoPE to Q/K and scatter K/V into the MLX KV pool.\n\n    Args:\n        q: Query tensor with shape `[num_tokens, num_qo_heads, head_dim]`.\n        k: Key tensor with shape `[num_tokens, num_kv_heads, head_dim]`.\n        v: Value tensor with shape `[num_tokens, num_kv_heads, head_dim]`.\n        positions: int32 positions with shape `[num_tokens]`.\n        slots: int32 KV-pool slots with shape `[num_tokens]`; values `< 0`\n            skip the pool write for that token.\n        k_pool: Existing K pool with shape `[pool_size, num_kv_heads, head_dim]`.\n        v_pool: Existing V pool with shape `[pool_size, num_kv_heads, head_dim]`.\n\n    Returns:\n        `(q_rot, k_rot, k_pool_new, v_pool_new)`.\n    \"\"\"\n    if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:\n        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(","sourceCodeStart":46,"sourceCodeEnd":82,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/aot/python/sgl_kernel/metal.py#L46-L82","documentation":"rope_pool_fused validates that the q, k, and v tensors passed for fused rotary embedding + KV pooling are 3-D with layout [num_tokens, num_heads, head_dim]. This error means at least one of q/k/v is not a 3-D tensor (e.g. a 2-D flattened tensor or a 4-D batched tensor). The check exists because the Metal kernel indexes tensors assuming exactly three dimensions.","triggerScenarios":"Calling rope_pool_fused(q, k, v, ...) where any of q/k/v has ndim != 3, e.g. passing a 4-D [batch, seq, heads, dim] attention-layout tensor or a 2-D [num_tokens, heads*dim] flattened projection output without reshaping.","commonSituations":"Adapter code that copies from CUDA-flavored flash-attention call sites where q/k/v are [batch, seq, heads, dim]; passing the raw QKV projection output without splitting/reshaping; hidden batch dimension from a dummy batch=1 wrapper.","solutions":["Reshape q/k/v to [num_tokens, num_qo_heads, head_dim] / [num_tokens, num_kv_heads, head_dim] before calling (e.g. q.view(num_tokens, num_qo_heads, head_dim))","If your tensors are [batch, seq, heads, dim], flatten batch and seq: q.reshape(-1, num_heads, head_dim)","Print q.ndim, k.ndim, v.ndim right before the call to identify the offending tensor"],"exampleFix":"# before\nq_out = metal.rope_pool_fused(q, k, v, ...)  # q is [1, seq, H, D]\n\n# after\nq = q.reshape(-1, num_qo_heads, head_dim)\nk = k.reshape(-1, num_kv_heads, head_dim)\nv = v.reshape(-1, num_kv_heads, head_dim)\nq_out = metal.rope_pool_fused(q, k, v, ...)","handlingStrategy":"validation","validationCode":"assert q.ndim == k.ndim == v.ndim == 3, (q.shape, k.shape, v.shape)","typeGuard":"def is_rope_qkv(q, k, v) -> bool:\n    return all(t.ndim == 3 and t.is_cuda is False or t.ndim == 3 for t in (q, k, v)) and q.ndim == 3","tryCatchPattern":null,"preventionTips":["Standardize on [num_tokens, heads, head_dim] layout at the projection site","Add shape asserts immediately after QKV split, before any kernel call"],"tags":["shape-validation","rope","metal","sgl-kernel","tensor-dims"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}