sgl-project/sglang · error · ValueError

Invalid fused KV projection shape: got {tuple(kv.shape)}, ex

Error message

Invalid fused KV projection shape: got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}.

What it means

The trailing dimension of the stacked fused KV tensor must equal exactly 2 * num_kv_heads * head_dim (K and V concatenated per layer). The code derives kv_size from num_kv_heads and head_dim and rejects any other trailing dim.

Source

Thrown at python/sglang/kernels/ops/speculative/fused_kv_materialize.py:154

    v_out: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fused RMSNorm + RoPE materialization for all layers."""
    if kv.ndim != 3:
        raise ValueError(
            "Invalid stacked fused KV projection shape: "
            f"got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2]."
        )

    total_ctx, n_layers, kv_dim = kv.shape
    if total_ctx == 0:
        empty = torch.empty(
            (n_layers, 0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device
        )
        return empty, empty

    kv_size = num_kv_heads * head_dim
    if kv_dim != kv_size * 2:
        raise ValueError(
            "Invalid fused KV projection shape: "
            f"got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}."
        )
    if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:
        raise ValueError(
            "Invalid fused KV rotary/head dim pair: "
            f"rotary_dim={rotary_dim}, head_dim={head_dim}."
        )
    if k_norm_weight.shape != (n_layers, head_dim):
        raise ValueError(
            "Invalid stacked k_norm_weight shape for fused KV materialization: "
            f"got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}."
        )
    if eps.shape != (n_layers,):
        raise ValueError(
            "Invalid stacked eps shape for fused KV materialization: "
            f"got {tuple(eps.shape)}, expected {(n_layers,)}."
        )

View on GitHub (pinned to 0132848349)

Solutions

Verify num_kv_heads and head_dim match the model config used to build the stacked projections.
Re-check that the fused qkv weight packs exactly K then V with total width 2*kv_size.
Print kv.shape vs expected num_kv_heads*head_dim*2 and reconcile the difference.

Example fix

// before
mat = FusedKVMaterializer(..., num_kv_heads=8, head_dim=128)  # kv trailing dim 4096
// after
mat = FusedKVMaterializer(..., num_kv_heads=8, head_dim=128)  # kv trailing dim must be 8*128*2=2048

Defensive patterns

Strategy: validation

Validate before calling

expected = num_kv_heads * head_dim * 2
assert kv.shape[-1] == expected, (kv.shape, expected)

Prevention

Derive num_kv_heads/head_dim from the same config used to build the projection weights.
Add a unit test asserting trailing-dim invariants.

When it happens

Trigger: Passing kv with last dim != num_kv_heads*head_dim*2, or calling with num_kv_heads/head_dim that don't match how the projection weights were fused.

Common situations: Model config mismatch: head_dim or num_kv_heads computed differently (e.g. derived from hidden_size/num_attention_heads) than the checkpoint's fused qkv weight layout.

Related errors

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