sgl-project/sglang · error · ValueError

k_pool has incompatible shape {k_pool.shape}

Error message

k_pool has incompatible shape {k_pool.shape}

What it means

The k_pool tensor must have trailing dimensions [num_kv_heads, head_dim]; its leading (pool_size) dimension is free. This error fires when k_pool.shape[1:] != (num_kv_heads, head_dim), i.e. the pool was allocated with head layout inconsistent with the declared KV configuration.

Source

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

    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}")
    if positions_shape != (q_shape[0],) or slots_shape != (q_shape[0],):
        raise ValueError("positions/slots must have one entry per token")
    if k_pool_shape[1:] != (num_kv_heads, head_dim):
        raise ValueError(f"k_pool has incompatible shape {k_pool.shape}")
    if v_pool_shape != k_pool_shape:
        raise ValueError(
            f"v_pool shape must match k_pool shape, got {v_pool.shape} vs {k_pool.shape}"
        )
    if q.dtype != k.dtype or q.dtype != v.dtype:
        raise ValueError("q/k/v dtypes must match")
    if k_pool.dtype != q.dtype or v_pool.dtype != q.dtype:
        raise ValueError("pool dtypes must match q/k/v dtype")

    return _metal.rope_pool_fused(
        q,
        k,
        v,
        positions,
        slots,
        k_pool,
        v_pool,
        head_dim,

View on GitHub (pinned to 0132848349)

Solutions

1. Reallocate k_pool as torch.empty(pool_size, num_kv_heads, head_dim, dtype=q.dtype)
2. If the pool is flat, reshape: k_pool = k_pool.view(pool_size, num_kv_heads, head_dim) (only if layout permits)
3. Verify num_kv_heads/head_dim passed to the call equal the values used at pool allocation

Example fix

# before
k_pool = torch.empty(pool_size, num_kv_heads * head_dim, dtype=dtype)

# after
k_pool = torch.empty(pool_size, num_kv_heads, head_dim, dtype=dtype)

assert k_pool.shape[1:] == (num_kv_heads, head_dim), (k_pool.shape, num_kv_heads, head_dim)

def pool_layout_ok(k_pool, num_kv_heads, head_dim) -> bool:
    return k_pool.ndim == 3 and tuple(k_pool.shape[1:]) == (num_kv_heads, head_dim)

When it happens

Trigger: Passing a pool allocated as [pool_size, num_kv_heads*head_dim], [pool_size, total_heads, head_dim] with MHA head count, or with head_dim of a different config.

Common situations: Reusing a cache allocated for another model or backend; allocating pools from a generic cache manager that flattens heads; config change (e.g. rope scaling variant) altering head_dim after pool allocation.

Related errors

• 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_pool shape must match k_pool shape, got {v_pool.shape} vs
• rope_pool_fused expects q/k/v to be 3-D

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