sgl-project/sglang · error · ValueError

k shape must be [num_tokens, num_kv_heads, head_dim], got {k

Error message

k shape must be [num_tokens, num_kv_heads, head_dim], got {k.shape}

What it means

rope_pool_fused validates that k has shape [num_tokens, num_kv_heads, head_dim], sharing num_tokens with q and using the declared KV head count. This error fires when k's shape disagrees — typically wrong number of KV heads or head_dim, or a different token count than q.

Source

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

        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}")
    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(

View on GitHub (pinned to 0132848349)

Solutions

1. Ensure k.shape == (q.shape[0], num_kv_heads, head_dim)
2. Derive num_kv_heads from the model config: num_kv_heads = k.shape[1] and pass that
3. Align token counts: k = k[: q.shape[0]] if q was sliced

Example fix

# before
num_kv_heads = config.num_attention_heads  # wrong for GQA

# after
num_kv_heads = config.num_key_value_heads
assert k.shape == (q.shape[0], num_kv_heads, head_dim)

assert k.shape == (q.shape[0], num_kv_heads, head_dim), (k.shape, num_kv_heads, head_dim)

def k_shape_ok(q, k, num_kv_heads, head_dim):
    return k.shape == (q.shape[0], num_kv_heads, head_dim)

When it happens

Trigger: Passing k with shape[1] != num_kv_heads or shape[2] != head_dim, or k.shape[0] != q.shape[0] (e.g. k trimmed/extended differently from q in a prefill path).

Common situations: GQA misconfiguration where num_kv_heads doesn't match the K projection width; slicing q and k with different token ranges; using MHA-style k (all heads) with GQA num_kv_heads from config.

Related errors

• q shape must be [num_tokens, num_qo_heads, head_dim], got {q
• rope_pool_fused expects q/k/v to be 3-D
• rope_pool_fused expects positions/slots to be 1-D
• rope_pool_fused expects pool tensors to be 3-D
• 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/a46e395c6e8da3f6. Report an issue: GitHub.