{"record":{"id":"f7b267cb130e0dd9","repo":"sgl-project/sglang","slug":"rope-pool-fused-expects-pool-tensors-to-be-3-d","errorCode":null,"errorMessage":"rope_pool_fused expects pool tensors to be 3-D","messagePattern":"rope_pool_fused expects pool tensors to be 3-D","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/aot/python/sgl_kernel/metal.py","lineNumber":68,"sourceCode":"    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(\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}\")","sourceCodeStart":50,"sourceCodeEnd":86,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/aot/python/sgl_kernel/metal.py#L50-L86","documentation":"The KV cache pool tensors k_pool and v_pool must be 3-D with layout [pool_size, num_kv_heads, head_dim]. This error means at least one pool tensor has a different rank (e.g. 4-D with an extra batch dim, or 2-D flattened). The kernel writes pooled keys/values by slot index into a flat 3-D pool.","triggerScenarios":"Passing a k_pool/v_pool of ndim != 3, such as a [layers, pool_size, heads, dim] layer-stacked cache, or a pool created with an extra leading dimension of size 1.","commonSituations":"Reusing a cache allocated for another backend (e.g. a 4-D paged cache); slicing a multi-layer pool but forgetting to index the layer dim; allocating the pool with torch.empty(pool_size, heads*dim) instead of (pool_size, heads, dim).","solutions":["Index or reshape the pool to exactly [pool_size, num_kv_heads, head_dim] (e.g. k_pool = cache[layer_idx])","Allocate pools as torch.empty(pool_size, num_kv_heads, head_dim, dtype=...)","Check k_pool.ndim and v_pool.ndim before calling"],"exampleFix":"# before\nk_pool = torch.empty(num_layers, pool_size, kv_heads, head_dim)\nmetal.rope_pool_fused(q, k, v, pos, slots, k_pool[0], v_pool[0], ...)  # wrong slicing left 4-D in some path\n\n# after\nk_pool_l = k_pool[layer_idx]\nv_pool_l = v_pool[layer_idx]\nmetal.rope_pool_fused(q, k, v, pos, slots, k_pool_l, v_pool_l, ...)","handlingStrategy":"validation","validationCode":"assert k_pool.ndim == v_pool.ndim == 3, (k_pool.shape, v_pool.shape)","typeGuard":"import torch\ndef is_pool3d(t: torch.Tensor) -> bool:\n    return t.ndim == 3","tryCatchPattern":null,"preventionTips":["Allocate pools with explicit (pool_size, num_kv_heads, head_dim)","Index the layer dimension of multi-layer caches before passing"],"tags":["shape-validation","kv-cache","metal","rope","sgl-kernel"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}