sgl-project/sglang · error · ValueError

rope_pool_fused expects q/k/v to be 3-D

Error message

rope_pool_fused expects q/k/v to be 3-D

What it means

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.

Source

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

    rope_base: float,
) -> tuple[mx.array, mx.array, mx.array, mx.array]:
    """Apply NeoX RoPE to Q/K and scatter K/V into the MLX KV pool.

    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(

View on GitHub (pinned to 0132848349)

Solutions

1. 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))
2. If your tensors are [batch, seq, heads, dim], flatten batch and seq: q.reshape(-1, num_heads, head_dim)
3. Print q.ndim, k.ndim, v.ndim right before the call to identify the offending tensor

Example fix

# before
q_out = metal.rope_pool_fused(q, k, v, ...)  # q is [1, seq, H, D]

# after
q = q.reshape(-1, num_qo_heads, head_dim)
k = k.reshape(-1, num_kv_heads, head_dim)
v = v.reshape(-1, num_kv_heads, head_dim)
q_out = metal.rope_pool_fused(q, k, v, ...)

assert q.ndim == k.ndim == v.ndim == 3, (q.shape, k.shape, v.shape)

def is_rope_qkv(q, k, v) -> bool:
    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

When it happens

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

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

Related errors

• rope_pool_fused expects positions/slots to be 1-D
• rope_pool_fused expects pool tensors to be 3-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
• v shape must match k shape, got {v.shape} vs {k.shape}

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