{"record":{"id":"4842247384a3636f","repo":"sgl-project/sglang","slug":"invalid-stacked-k-norm-weight-shape-for-fused-kv-m","errorCode":null,"errorMessage":"Invalid stacked k_norm_weight shape for fused KV materialization: got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}.","messagePattern":"Invalid stacked k_norm_weight shape for fused KV materialization: got (.+?), expected (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/speculative/fused_kv_materialize.py","lineNumber":164,"sourceCode":"    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:\n        raise ValueError(\n            \"Invalid fused KV rotary/head dim pair: \"\n            f\"rotary_dim={rotary_dim}, head_dim={head_dim}.\"\n        )\n    if k_norm_weight.shape != (n_layers, head_dim):\n        raise ValueError(\n            \"Invalid stacked k_norm_weight shape for fused KV materialization: \"\n            f\"got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}.\"\n        )\n    if eps.shape != (n_layers,):\n        raise ValueError(\n            \"Invalid stacked eps shape for fused KV materialization: \"\n            f\"got {tuple(eps.shape)}, expected {(n_layers,)}.\"\n        )\n\n    half_rotary_dim = rotary_dim // 2\n    BLOCK_HD = triton.next_power_of_2(head_dim)\n\n    if positions.device != kv.device:\n        positions = positions.to(device=kv.device, dtype=torch.int64)\n    elif positions.dtype != torch.int64:\n        positions = positions.to(torch.int64)\n\n    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)","sourceCodeStart":146,"sourceCodeEnd":182,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/speculative/fused_kv_materialize.py#L146-L182","documentation":"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.","triggerScenarios":"Passing a single [head_dim] norm weight (not stacked across layers) or weights shaped [n_layers, 1, head_dim].","commonSituations":"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.","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."],"exampleFix":"// before\nk_norm = layers[0].k_norm.weight  # [head_dim]\n// after\nk_norm = torch.stack([l.k_norm.weight for l in layers])  # [n_layers, head_dim]","handlingStrategy":"validation","validationCode":"assert k_norm_weight.shape == (n_layers, head_dim)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use torch.stack over per-layer weights instead of passing one layer's weight.","Let the materializer class collect weights automatically."],"tags":["shape-validation","rmsnorm","speculative-decoding"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}