{"record":{"id":"ef2f4d83f9f03e5e","repo":"sgl-project/sglang","slug":"invalid-stacked-fused-kv-projection-shape-got-tu","errorCode":null,"errorMessage":"Invalid stacked fused KV projection shape: got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2].","messagePattern":"Invalid stacked fused KV projection shape: got (.+?), expected 3D \\[total_ctx, n_layers, kv_size\\*2\\]\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/speculative/fused_kv_materialize.py","lineNumber":140,"sourceCode":"    mask_pass = (offs >= rotary_dim) & (offs < head_dim)\n    tl.store(k_write + offs, k_normed.to(v_raw.dtype), mask=mask_pass)\n\n\ndef _fused_norm_rope_stacked(\n    kv: torch.Tensor,  # [total_ctx, n_layers, kv_size*2]\n    k_norm_weight: torch.Tensor,  # [n_layers, head_dim]\n    eps: torch.Tensor,  # [n_layers]\n    cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim]\n    positions: torch.Tensor,  # [total_ctx]\n    num_kv_heads: int,\n    head_dim: int,\n    rotary_dim: int,\n    k_out: Optional[torch.Tensor] = None,\n    v_out: Optional[torch.Tensor] = None,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Fused RMSNorm + RoPE materialization for all layers.\"\"\"\n    if kv.ndim != 3:\n        raise ValueError(\n            \"Invalid stacked fused KV projection shape: \"\n            f\"got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2].\"\n        )\n\n    total_ctx, n_layers, kv_dim = kv.shape\n    if total_ctx == 0:\n        empty = torch.empty(\n            (n_layers, 0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device\n        )\n        return empty, empty\n\n    kv_size = num_kv_heads * head_dim\n    if kv_dim != kv_size * 2:\n        raise ValueError(\n            \"Invalid fused KV projection shape: \"\n            f\"got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}.\"\n        )\n    if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0:","sourceCodeStart":122,"sourceCodeEnd":158,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/speculative/fused_kv_materialize.py#L122-L158","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"// before\nk, v = mat.materialize(kv[layer0_only], positions)  # 2D slice\n// after\nk, v = mat.materialize(stacked_kv, positions)  # [total_ctx, n_layers, kv_size*2]","handlingStrategy":"validation","validationCode":"assert kv.ndim == 3, f'expected 3D stacked KV, got {kv.shape}'","typeGuard":"def is_stacked_kv(t: torch.Tensor) -> bool:\n    return t.ndim == 3","tryCatchPattern":null,"preventionTips":["Always build the stacked tensor with the provided stacking helper rather than manual concatenation.","Assert ndim==3 in tests covering the fused KV path."],"tags":["shape-validation","speculative-decoding","fused-kernel"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}