sgl-project/sglang · error · ValueError

Invalid stacked fused KV projection shape: got {tuple(kv.sha

Error message

Invalid stacked fused KV projection shape: got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2].

What it means

The stacked fused KV projection tensor passed to the fused RMSNorm+RoPE materialization path must be 3D with layout [total_ctx, n_layers, kv_size*2]. This check fires when the tensor's rank is anything other than 3, e.g. a per-layer 2D projection was passed where a stacked all-layers tensor is required.

Source

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

    mask_pass = (offs >= rotary_dim) & (offs < head_dim)
    tl.store(k_write + offs, k_normed.to(v_raw.dtype), mask=mask_pass)


def _fused_norm_rope_stacked(
    kv: torch.Tensor,  # [total_ctx, n_layers, kv_size*2]
    k_norm_weight: torch.Tensor,  # [n_layers, head_dim]
    eps: torch.Tensor,  # [n_layers]
    cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim]
    positions: torch.Tensor,  # [total_ctx]
    num_kv_heads: int,
    head_dim: int,
    rotary_dim: int,
    k_out: Optional[torch.Tensor] = None,
    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:

View on GitHub (pinned to 0132848349)

Solutions

Check kv.ndim and kv.shape before calling; the tensor must be [total_ctx, n_layers, num_kv_heads*head_dim*2].
Ensure the model produces a stacked fused KV projection across all layers (use the helper that builds the stacked tensor).
If you only have per-layer projections, stack them along dim=1 before calling materialize.

Example fix

// before
k, v = mat.materialize(kv[layer0_only], positions)  # 2D slice
// after
k, v = mat.materialize(stacked_kv, positions)  # [total_ctx, n_layers, kv_size*2]

Defensive patterns

Strategy: validation

Validate before calling

assert kv.ndim == 3, f'expected 3D stacked KV, got {kv.shape}'

Type guard

def is_stacked_kv(t: torch.Tensor) -> bool:
    return t.ndim == 3

Prevention

Always build the stacked tensor with the provided stacking helper rather than manual concatenation.
Assert ndim==3 in tests covering the fused KV path.

When it happens

Trigger: Calling materialize() (which calls _fused_norm_rope_stacked) with a kv tensor that has ndim != 3 — e.g. a single layer's [tokens, 2*kv_size] projection, or a batched 4D tensor.

Common situations: Integrating the fused KV speculative path with a model whose qkv projections are not stacked across layers, or accidentally slicing the stacked tensor before passing it in.

Related errors

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