From 024639a3722be0c7818237368ff6c91e090838e5 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 21 Jul 2026 17:46:13 -0700 Subject: [PATCH] [Perf] Stack dspark dense draft per-layer ctx KV projection into one GEMM (#31986) --- python/sglang/srt/models/dspark.py | 103 +++++++++- .../test_dspark_stacked_ctx_kv_parity.py | 177 ++++++++++++++++++ 2 files changed, 274 insertions(+), 6 deletions(-) create mode 100644 test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index 0ffc22b3b..a1ca3f145 100644 --- a/python/sglang/srt/models/dspark.py +++ b/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): diff --git a/test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py b/test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py new file mode 100644 index 000000000..951870fe7 --- /dev/null +++ b/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()