sgl-project/sglang · error · ValueError

num_kv_heads mismatch across layers for fused KV path: expec

Error message

num_kv_heads mismatch across layers for fused KV path: expected {self.num_kv_heads}, got {int(attn.num_kv_heads)} at layer {layer_id}.

What it means

All decoder layers must have the same num_kv_heads for the stacked fused KV path, since one Triton launch handles all layers with a single head-count grid. The check compares each layer's attn.num_kv_heads against the first layer's.

Source

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

        )
        self._reserved_rope_cache_len = int(
            getattr(self.rotary_emb, "cos_sin_cache", torch.empty((0,))).shape[0]
        )
        self._mm_out_supported = True
        self._workspace_capacity = 0
        self._workspace_dtype: Optional[torch.dtype] = None
        self._proj_workspace: Optional[torch.Tensor] = None
        self._k_workspace: Optional[torch.Tensor] = None
        self._v_workspace: Optional[torch.Tensor] = None

        kv_weights = []
        k_norm_weights = []
        eps_values = []

        for layer_id, layer in enumerate(layers):
            attn = layer.self_attn
            if int(attn.num_kv_heads) != self.num_kv_heads:
                raise ValueError(
                    "num_kv_heads mismatch across layers for fused KV path: "
                    f"expected {self.num_kv_heads}, got {int(attn.num_kv_heads)} at layer {layer_id}."
                )
            if int(attn.head_dim) != self.head_dim:
                raise ValueError(
                    "head_dim mismatch across layers for fused KV path: "
                    f"expected {self.head_dim}, got {int(attn.head_dim)} at layer {layer_id}."
                )
            layer_rotary_dim = int(
                getattr(attn.rotary_emb, "rotary_dim", self.head_dim)
            )
            layer_is_neox = bool(getattr(attn.rotary_emb, "is_neox_style", True))
            if (
                layer_rotary_dim != self.rotary_dim
                or layer_is_neox != self.is_neox_style
            ):
                raise ValueError(
                    "RoPE config mismatch across layers for fused KV path: "

View on GitHub (pinned to 0132848349)

Solutions

  1. Only enable the fused KV path on models with uniform num_kv_heads across layers.
  2. Verify each layer's config in the checkpoint; fix any wrong per-layer overrides.
  3. Fall back to the non-fused per-layer materialization for heterogeneous models.
Defensive patterns

Strategy: validation

Validate before calling

heads = {int(l.self_attn.num_kv_heads) for l in layers}
assert len(heads) == 1

Type guard

def uniform_kv_heads(layers) -> bool:
    return len({int(l.self_attn.num_kv_heads) for l in layers}) == 1

Prevention

When it happens

Trigger: A model where some attention layers use different num_kv_heads (e.g. alternating GQA group sizes or layer-wise heterogeneous attention).

Common situations: New hybrid/checkpoint architectures with per-layer attention config arrays; passing a subset of layers with mismatched configs.

Related errors


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