sgl-project/sglang · error · ValueError

pool dtypes must match q/k/v dtype

Error message

pool dtypes must match q/k/v dtype

What it means

The Metal rope_pool_fused wrapper requires the k_pool and v_pool tensors to have the exact same dtype as q/k/v. The fused kernel writes RoPE-rotated and pooled values with a single compiled element type, so mismatched pool dtypes are rejected.

Source

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

        )
    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,
        num_qo_heads,
        num_kv_heads,
        float(rope_base),
    )

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate k_pool/v_pool with torch.empty(..., dtype=q.dtype)
2. Cast pools to q.dtype before the call: k_pool=k_pool.to(q.dtype)
3. If float32 pooling is required, use a kernel variant that supports it instead of this fused op

Example fix

// before
k_pool = torch.zeros(n, dtype=torch.float32)
rope_pool_fused(q, k, v, k_pool, v_pool, ...)
// after
k_pool = torch.zeros(n, dtype=q.dtype)
rope_pool_fused(q, k, v, k_pool, v_pool, ...)

assert k_pool.dtype == v_pool.dtype == q.dtype

except ValueError: k_pool = k_pool.to(q.dtype); v_pool = v_pool.to(q.dtype); retry

When it happens

Trigger: Calling rope_pool_fused where k_pool.dtype or v_pool.dtype differs from q.dtype, e.g. bf16 q with fp32 k_pool, or fp16 v_pool with bf16 q.

Common situations: Pre-allocated KV pools in float32 for accumulation while the model runs in bf16/fp16; reusing pools allocated by a different backend; partial casts applied to pools but not q/k/v.

Related errors

• q/k/v dtypes must match
• 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
• q shape must be [num_tokens, num_qo_heads, head_dim], got {q

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