sgl-project/sglang · error · ValueError

Invalid stacked k_norm_weight shape for fused KV materializa

Error message

Invalid stacked k_norm_weight shape for fused KV materialization: got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}.

What it means

k_norm_weight for the stacked fused KV materialization must be a 2D tensor [n_layers, head_dim] — one RMSNorm weight vector per layer. Any other shape is rejected.

Source

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

    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,)}."
        )

    half_rotary_dim = rotary_dim // 2
    BLOCK_HD = triton.next_power_of_2(head_dim)

    if positions.device != kv.device:
        positions = positions.to(device=kv.device, dtype=torch.int64)
    elif positions.dtype != torch.int64:
        positions = positions.to(torch.int64)

    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)

View on GitHub (pinned to 0132848349)

Solutions

Stack per-layer k_norm weights: torch.stack([layer.k_norm.weight for layer in layers]) giving [n_layers, head_dim].
If the model has no k_norm, pass ones of shape (n_layers, head_dim).
Verify n_layers matches the kv tensor's dim 1.

Example fix

// before
k_norm = layers[0].k_norm.weight  # [head_dim]
// after
k_norm = torch.stack([l.k_norm.weight for l in layers])  # [n_layers, head_dim]

Defensive patterns

Strategy: validation

Validate before calling

assert k_norm_weight.shape == (n_layers, head_dim)

Prevention

Use torch.stack over per-layer weights instead of passing one layer's weight.
Let the materializer class collect weights automatically.

When it happens

Trigger: Passing a single [head_dim] norm weight (not stacked across layers) or weights shaped [n_layers, 1, head_dim].

Common situations: Model without per-layer k_norm (e.g. Qwen-style q/k norm absent) where a broadcast weight was fabricated, or forgetting to stack per-layer norm weights when building inputs manually.

Related errors

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