[Perf] Stack dspark dense draft per-layer ctx KV projection into one GEMM (#31986)

python/sglang/srt/models/dspark.py

@@ -4,11 +4,13 @@ import logging
 from typing import Callable, Iterable, Optional, Tuple
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.dflash import DFlashDraftModel
+from sglang.srt.speculative.dflash_utils import can_dflash_slice_qkv_weight
 from sglang.srt.speculative.dspark_components.dspark_config import (
     parse_dspark_draft_config,
 )
@@ -464,6 +466,42 @@ class DSparkDraftMixin:
                 f"or disable the confidence head (enable_confidence_head=False)."
             )
 
+    def _stacked_ctx_kv_params(self) -> Optional[dict]:
+        """Stack every layer's KV projection into one weight (exact: the input
+        hidden is shared, so concatenating output columns is equivalent).
+        Cached; None (per-layer fallback) when a QKV weight cannot be sliced
+        (quantized) or layers disagree on norm epsilon / bias presence.
+        """
+        cached = getattr(self, "_stacked_ctx_kv_cache", False)
+        if cached is not False:
+            return cached
+        weights, biases, k_norm_weights = [], [], []
+        eps = None
+        for layer in self.layers:
+            attn = layer.self_attn
+            can_slice, _ = can_dflash_slice_qkv_weight(attn.qkv_proj)
+            if not can_slice or eps not in (None, attn.k_norm.variance_epsilon):
+                self._stacked_ctx_kv_cache = None
+                return None
+            eps = attn.k_norm.variance_epsilon
+            kv_slice = slice(attn.q_size, attn.q_size + 2 * attn.kv_size)
+            weights.append(attn.qkv_proj.weight[kv_slice])
+            biases.append(
+                attn.qkv_proj.bias[kv_slice] if attn.qkv_proj.bias is not None else None
+            )
+            k_norm_weights.append(attn.k_norm.weight)
+        has_bias = [b is not None for b in biases]
+        if any(has_bias) and not all(has_bias):
+            self._stacked_ctx_kv_cache = None
+            return None
+        self._stacked_ctx_kv_cache = {
+            "weight": torch.cat(weights, dim=0),
+            "bias": torch.cat(biases, dim=0) if all(has_bias) else None,
+            "k_norm_weight": torch.stack(k_norm_weights, dim=0).float(),
+            "eps": eps,
+        }
+        return self._stacked_ctx_kv_cache
+
     def write_target_hidden_kv(
         self,
         *,
@@ -475,13 +513,22 @@ class DSparkDraftMixin:
         commit_lens: Optional[torch.Tensor] = None,
     ) -> None:
         ctx_hidden = self.project_target_hidden(target_hidden)
-        for layer in self.layers:
+        stacked = self._stacked_ctx_kv_params()
+        if stacked is not None:
+            k_all, v_all = self._project_ctx_kv_stacked(
+                ctx_hidden=ctx_hidden, positions=positions, stacked=stacked
+            )
+        for i, layer in enumerate(self.layers):
             attn = layer.self_attn
-            k, v = attn.kv_proj_only(ctx_hidden)
-            k = attn.apply_k_norm(k)
-            k = attn.apply_k_rope(positions, k)
-            k = k.view(-1, attn.num_kv_heads, attn.head_dim)
-            v = v.view(-1, attn.num_kv_heads, attn.head_dim)
+            if stacked is not None:
+                k = k_all[i]
+                v = v_all[i]
+            else:
+                k, v = attn.kv_proj_only(ctx_hidden)
+                k = attn.apply_k_norm(k)
+                k = attn.apply_k_rope(positions, k)
+                k = k.view(-1, attn.num_kv_heads, attn.head_dim)
+                v = v.view(-1, attn.num_kv_heads, attn.head_dim)
             if cache_loc_2d is not None and commit_lens is not None:
                 pool.set_kv_buffer_prefix_valid(
                     attn.attn,
@@ -502,6 +549,50 @@ class DSparkDraftMixin:
                     attn.attn.v_scale,
                 )
 
+    def _project_ctx_kv_stacked(
+        self,
+        *,
+        ctx_hidden: torch.Tensor,
+        positions: torch.Tensor,
+        stacked: dict,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        attn0 = self.layers[0].self_attn
+        num_layers = len(self.layers)
+        kv_size = attn0.kv_size
+        head_dim = attn0.head_dim
+        num_kv_heads = attn0.num_kv_heads
+        tokens = ctx_hidden.shape[0]
+
+        kv_all = F.linear(ctx_hidden, stacked["weight"], stacked["bias"])
+        kv_all = kv_all.view(tokens, num_layers, 2, kv_size)
+        # Batched per-head k-norm across layers (fp32 variance + weight, cast back).
+        k32 = (
+            kv_all[:, :, 0, :]
+            .reshape(tokens, num_layers, num_kv_heads, head_dim)
+            .to(torch.float32)
+        )
+        variance = k32.pow(2).mean(dim=-1, keepdim=True)
+        k32 = k32 * torch.rsqrt(variance + stacked["eps"])
+        k32 = k32 * stacked["k_norm_weight"].view(1, num_layers, 1, head_dim)
+        k_all = k32.to(ctx_hidden.dtype)
+        # One RoPE over all layers' heads (shared rotary params + positions).
+        k_flat = k_all.reshape(tokens, num_layers * kv_size)
+        dummy_q = k_flat.new_empty(k_flat.shape)
+        _, k_flat = attn0.rotary_emb(positions, dummy_q, k_flat)
+        # [layers, tokens, heads, dim]: per-layer slices are contiguous views.
+        k_all = (
+            k_flat.view(tokens, num_layers, num_kv_heads, head_dim)
+            .permute(1, 0, 2, 3)
+            .contiguous()
+        )
+        v_all = (
+            kv_all[:, :, 1, :]
+            .view(tokens, num_layers, num_kv_heads, head_dim)
+            .permute(1, 0, 2, 3)
+            .contiguous()
+        )
+        return k_all, v_all
+
 
 class DSparkDraftModel(DSparkDraftMixin, DFlashDraftModel):
 

test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py

@@ -0,0 +1,177 @@
+"""Parity: dspark stacked ctx-KV write vs the per-layer loop.
+
+Accuracy tests cannot catch a broken ctx-KV write (spec decoding stays correct
+regardless of draft KV; only accept length drops), so compare the two paths
+directly and check the fallbacks return None.
+"""
+
+import types
+import unittest
+
+import torch
+
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.layers.rotary_embedding import get_rope
+from sglang.srt.models.dflash import DFlashAttention
+from sglang.srt.models.dspark import DSparkDraftMixin
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
+
+DEVICE = torch.device("cuda")
+HEAD_DIM = 64
+NUM_KV_HEADS = 2
+NUM_Q_HEADS = 4
+HIDDEN = 128
+EPS = 1e-6
+
+
+class _MockQKV:
+    """Fused-QKV stand-in that satisfies can_dflash_slice_qkv_weight."""
+
+    def __init__(self, weight, bias, quantized=False):
+        self.quant_method = object() if quantized else UnquantizedLinearMethod()
+        self.weight = weight
+        self.bias = bias
+
+
+def _make_attn(rope, *, eps=EPS, has_bias=False, quantized=False, g=None):
+    attn = types.SimpleNamespace()
+    attn.num_kv_heads = NUM_KV_HEADS
+    attn.head_dim = HEAD_DIM
+    attn.q_size = NUM_Q_HEADS * HEAD_DIM
+    attn.kv_size = NUM_KV_HEADS * HEAD_DIM
+    out = attn.q_size + 2 * attn.kv_size
+    weight = torch.randn(out, HIDDEN, device=DEVICE, dtype=torch.float32, generator=g)
+    bias = (
+        torch.randn(out, device=DEVICE, dtype=torch.float32, generator=g)
+        if has_bias
+        else None
+    )
+    attn.qkv_proj = _MockQKV(weight, bias, quantized=quantized)
+    k_norm = RMSNorm(HEAD_DIM, eps=eps).to(DEVICE)
+    with torch.no_grad():
+        # Distinct per layer so a wrong layer order fails parity.
+        k_norm.weight.copy_(torch.randn(HEAD_DIM, device=DEVICE, generator=g))
+    attn.k_norm = k_norm
+    attn.rotary_emb = rope
+    for name in ("kv_proj_only", "apply_k_norm", "apply_k_rope"):
+        setattr(attn, name, types.MethodType(getattr(DFlashAttention, name), attn))
+    return attn
+
+
+def _make_model(rope, num_layers, **kw):
+    layers = [
+        types.SimpleNamespace(self_attn=_make_attn(rope, **kw))
+        for _ in range(num_layers)
+    ]
+    model = types.SimpleNamespace(layers=layers)
+    for name in ("_stacked_ctx_kv_params", "_project_ctx_kv_stacked"):
+        setattr(model, name, types.MethodType(getattr(DSparkDraftMixin, name), model))
+    return model
+
+
+def _per_layer_reference(model, ctx_hidden, positions):
+    ks, vs = [], []
+    for layer in model.layers:
+        attn = layer.self_attn
+        k, v = attn.kv_proj_only(ctx_hidden)
+        k = attn.apply_k_norm(k)
+        k = attn.apply_k_rope(positions, k)
+        ks.append(k.view(-1, attn.num_kv_heads, attn.head_dim))
+        vs.append(v.view(-1, attn.num_kv_heads, attn.head_dim))
+    return ks, vs
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+class TestDSparkStackedCtxKvParity(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+        self.rope = get_rope(
+            HEAD_DIM,
+            rotary_dim=HEAD_DIM,
+            max_position=4096,
+            base=10000.0,
+            is_neox_style=True,
+        ).to(DEVICE)
+
+    def _check_parity(self, *, num_layers=4, tokens=5, has_bias=False, dtype):
+        g = torch.Generator(device=DEVICE).manual_seed(0)
+        model = _make_model(self.rope, num_layers, has_bias=has_bias, g=g)
+        for layer in model.layers:
+            attn = layer.self_attn
+            attn.qkv_proj.weight = attn.qkv_proj.weight.to(dtype)
+            if attn.qkv_proj.bias is not None:
+                attn.qkv_proj.bias = attn.qkv_proj.bias.to(dtype)
+            attn.k_norm.to(dtype)
+        ctx_hidden = torch.randn(
+            tokens, HIDDEN, device=DEVICE, dtype=dtype, generator=g
+        )
+        positions = torch.arange(tokens, device=DEVICE)
+
+        ref_k, ref_v = _per_layer_reference(model, ctx_hidden, positions)
+        stacked = model._stacked_ctx_kv_params()
+        self.assertIsNotNone(stacked)
+        k_all, v_all = model._project_ctx_kv_stacked(
+            ctx_hidden=ctx_hidden, positions=positions, stacked=stacked
+        )
+
+        # Tol covers the fused-kernel-vs-manual-fp32 RMSNorm rounding; O(1)
+        # wiring errors fail either way. (Fused rmsnorm has no fp32 dispatch.)
+        rtol, atol = {torch.float16: (5e-3, 5e-3), torch.bfloat16: (2e-2, 2e-2)}[dtype]
+        for i in range(num_layers):
+            torch.testing.assert_close(k_all[i], ref_k[i], rtol=rtol, atol=atol)
+            torch.testing.assert_close(v_all[i], ref_v[i], rtol=rtol, atol=atol)
+
+    def test_parity_fp16(self):
+        self._check_parity(dtype=torch.float16)
+
+    def test_parity_bf16(self):
+        self._check_parity(dtype=torch.bfloat16)
+
+    def test_parity_with_bias(self):
+        self._check_parity(dtype=torch.float16, has_bias=True)
+
+    def test_fallback_quantized_layer(self):
+        g = torch.Generator(device=DEVICE).manual_seed(0)
+        model = _make_model(self.rope, 3, g=g)
+        model.layers[1].self_attn.qkv_proj.quant_method = object()
+        self.assertIsNone(model._stacked_ctx_kv_params())
+
+    def test_fallback_eps_mismatch(self):
+        g = torch.Generator(device=DEVICE).manual_seed(0)
+        model = types.SimpleNamespace(
+            layers=[
+                types.SimpleNamespace(self_attn=_make_attn(self.rope, eps=1e-6, g=g)),
+                types.SimpleNamespace(self_attn=_make_attn(self.rope, eps=1e-5, g=g)),
+            ]
+        )
+        model._stacked_ctx_kv_params = types.MethodType(
+            DSparkDraftMixin._stacked_ctx_kv_params, model
+        )
+        self.assertIsNone(model._stacked_ctx_kv_params())
+
+    def test_fallback_inconsistent_bias(self):
+        g = torch.Generator(device=DEVICE).manual_seed(0)
+        model = types.SimpleNamespace(
+            layers=[
+                types.SimpleNamespace(
+                    self_attn=_make_attn(self.rope, has_bias=True, g=g)
+                ),
+                types.SimpleNamespace(
+                    self_attn=_make_attn(self.rope, has_bias=False, g=g)
+                ),
+            ]
+        )
+        model._stacked_ctx_kv_params = types.MethodType(
+            DSparkDraftMixin._stacked_ctx_kv_params, model
+        )
+        self.assertIsNone(model._stacked_ctx_kv_params())
+
+
+if __name__ == "__main__":
+    unittest.main()