sgl-project/sglang · error · ValueError

rope_pool_fused expects pool tensors to be 3-D

Error message

rope_pool_fused expects pool tensors to be 3-D

What it means

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.

Source

Thrown at python/sglang/kernels/aot/python/sgl_kernel/metal.py:68

    Args:
        q: Query tensor with shape `[num_tokens, num_qo_heads, head_dim]`.
        k: Key tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
        v: Value tensor with shape `[num_tokens, num_kv_heads, head_dim]`.
        positions: int32 positions with shape `[num_tokens]`.
        slots: int32 KV-pool slots with shape `[num_tokens]`; values `< 0`
            skip the pool write for that token.
        k_pool: Existing K pool with shape `[pool_size, num_kv_heads, head_dim]`.
        v_pool: Existing V pool with shape `[pool_size, num_kv_heads, head_dim]`.

    Returns:
        `(q_rot, k_rot, k_pool_new, v_pool_new)`.
    """
    if q.ndim != 3 or k.ndim != 3 or v.ndim != 3:
        raise ValueError("rope_pool_fused expects q/k/v to be 3-D")
    if positions.ndim != 1 or slots.ndim != 1:
        raise ValueError("rope_pool_fused expects positions/slots to be 1-D")
    if k_pool.ndim != 3 or v_pool.ndim != 3:
        raise ValueError("rope_pool_fused expects pool tensors to be 3-D")
    q_shape = tuple(q.shape)
    k_shape = tuple(k.shape)
    v_shape = tuple(v.shape)
    positions_shape = tuple(positions.shape)
    slots_shape = tuple(slots.shape)
    k_pool_shape = tuple(k_pool.shape)
    v_pool_shape = tuple(v_pool.shape)

    if q_shape != (q_shape[0], num_qo_heads, head_dim):
        raise ValueError(
            "q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}"
        )
    if k_shape != (q_shape[0], num_kv_heads, head_dim):
        raise ValueError(
            "k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}"
        )
    if v_shape != k_shape:
        raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}")

View on GitHub (pinned to 0132848349)

Solutions

1. Index or reshape the pool to exactly [pool_size, num_kv_heads, head_dim] (e.g. k_pool = cache[layer_idx])
2. Allocate pools as torch.empty(pool_size, num_kv_heads, head_dim, dtype=...)
3. Check k_pool.ndim and v_pool.ndim before calling

Example fix

# before
k_pool = torch.empty(num_layers, pool_size, kv_heads, head_dim)
metal.rope_pool_fused(q, k, v, pos, slots, k_pool[0], v_pool[0], ...)  # wrong slicing left 4-D in some path

# after
k_pool_l = k_pool[layer_idx]
v_pool_l = v_pool[layer_idx]
metal.rope_pool_fused(q, k, v, pos, slots, k_pool_l, v_pool_l, ...)

assert k_pool.ndim == v_pool.ndim == 3, (k_pool.shape, v_pool.shape)

import torch
def is_pool3d(t: torch.Tensor) -> bool:
    return t.ndim == 3

When it happens

Trigger: 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.

Common situations: 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).

Related errors

• v_pool shape must match k_pool shape, got {v_pool.shape} vs
• rope_pool_fused expects q/k/v to be 3-D
• rope_pool_fused expects positions/slots to be 1-D
• q shape must be [num_tokens, num_qo_heads, head_dim], got {q
• k shape must be [num_tokens, num_kv_heads, head_dim], got {k

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/f7b267cb130e0dd9. Report an issue: GitHub.