[Perf][DSpark] Stack DeepSeek V4 context WKV projections (#54674)
Signed-off-by: liuyao0322 <yaoliu548926@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
leo committed
Sep 17, 2026 at 15:04 UTC
b3079e6e46a21109a8641c6d48641453a9c7cb86
2 files changed
+119
-6
tests/models/test_dspark_mla.py
+73
@@ -11,6 +11,7 @@ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
11
from vllm.model_executor.layers.logits_processor import LogitsProcessor
12
from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
13
from vllm.model_executor.models.registry import ModelRegistry
14
+from vllm.models.deepseek_v4.nvidia import dspark as dsv4_dspark
15
from vllm.models.kimi_k3.nvidia import dspark_mla
16
from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel
17
@@ -262,3 +263,75 @@ def test_v41_dspark_loads_linear_scales(
263
assert loaded == {runtime_name}
264
assert shards == [() if shard_id is None else (shard_id,)]
265
torch.testing.assert_close(param, checkpoint_scale)
266
+
267
+
268
+def test_dsv4_context_wkv_weights_are_duplicated_by_draft_layer():
269
+ weights = [
270
+ ("mtp.0.attn.wkv.weight", torch.arange(4)),
271
+ ("mtp.1.attn.wq_a.weight", torch.arange(3)),
272
+ ("mtp.2.attn.wkv.scale", torch.tensor(0.5)),
273
+ ("mtp.3.attn.wkv.weight", torch.arange(2)),
274
+ ]
275
+
276
+ duplicated = list(dsv4_dspark._duplicate_context_wkv_weights(weights, 3))
277
+
278
+ assert [name for name, _ in duplicated] == [
279
+ "mtp.0.attn.wkv.weight",
280
+ "context_wkv_proj.weight",
281
+ "mtp.1.attn.wq_a.weight",
282
+ "mtp.2.attn.wkv.scale",
283
+ "context_wkv_proj.scale",
284
+ "mtp.3.attn.wkv.weight",
285
+ ]
286
+ assert duplicated[1][1].shard_id == 0
287
+ assert duplicated[4][1].shard_id == 2
288
+ assert duplicated[0][1].data_ptr() == duplicated[1][1].data_ptr()
289
+ assert duplicated[3][1].data_ptr() == duplicated[4][1].data_ptr()
290
+
291
+
292
+def test_dsv4_context_kv_uses_one_stacked_wkv_projection(monkeypatch):
293
+ calls = []
294
+ stacked_output = torch.arange(24, dtype=torch.float32).view(2, 12)
295
+
296
+ class StackedProjection:
297
+ def __init__(self):
298
+ self.calls = 0
299
+
300
+ def __call__(self, main_x):
301
+ self.calls += 1
302
+ assert main_x.shape == (2, 5)
303
+ return stacked_output
304
+
305
+ projection = StackedProjection()
306
+ layers = [
307
+ SimpleNamespace(attn=SimpleNamespace(kv_norm=lambda kv, offset=i: kv + offset))
308
+ for i in range(3)
309
+ ]
310
+ model = SimpleNamespace(
311
+ config=SimpleNamespace(head_dim=4),
312
+ context_wkv_proj=projection,
313
+ layers=layers,
314
+ num_dspark_layers=3,
315
+ )
316
+ slot_mappings = [torch.tensor([0, 1]), None, torch.tensor([4, 5])]
317
+ monkeypatch.setattr(
318
+ dsv4_dspark,
319
+ "_insert_context_kv",
320
+ lambda attn, kv, positions, slots: calls.append(
321
+ (attn, kv.clone(), positions, slots)
322
+ ),
323
+ )
324
+
325
+ dsv4_dspark.DSparkDeepseekV4Model.precompute_and_store_context_kv(
326
+ model,
327
+ torch.zeros(2, 5),
328
+ torch.tensor([7, 8]),
329
+ slot_mappings,
330
+ )
331
+
332
+ assert projection.calls == 1
333
+ assert len(calls) == 2
334
+ assert torch.equal(calls[0][1], stacked_output.view(2, 3, 4)[:, 0])
335
+ assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2)
336
+ assert calls[0][3] is slot_mappings[0]
337
+ assert calls[1][3] is slot_mappings[2]
vllm/models/deepseek_v4/nvidia/dspark.py
+46
-6
@@ -31,7 +31,10 @@ from vllm.model_executor.layers.fused_moe import (
31
fused_moe_make_expert_params_mapping,
32
)
33
from vllm.model_executor.layers.layernorm import RMSNorm
34
-from vllm.model_executor.layers.linear import ReplicatedLinear
34
+from vllm.model_executor.layers.linear import (
35
+ MergedColumnParallelLinear,
36
+ ReplicatedLinear,
37
+)
38
from vllm.model_executor.layers.logits_processor import LogitsProcessor
39
from vllm.model_executor.layers.vocab_parallel_embedding import (
40
ParallelLMHead,
@@ -61,6 +64,24 @@ logger = init_logger(__name__)
64
# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
65
# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
66
_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
67
+_CONTEXT_WKV_RE = re.compile(r"^mtp\.(\d+)\.attn\.wkv\.(.+)$")
68
+
69
+
70
+def _duplicate_context_wkv_weights(
71
+ weights: Iterable[tuple[str, torch.Tensor]], num_layers: int
72
+) -> Iterable[tuple[str, torch.Tensor]]:
73
+ """Load every draft layer's WKV into the cross-layer projection."""
74
+ for name, weight in weights:
75
+ yield name, weight
76
+ match = _CONTEXT_WKV_RE.fullmatch(name)
77
+ if match is None:
78
+ continue
79
+ layer_idx = int(match.group(1))
80
+ if layer_idx >= num_layers:
81
+ continue
82
+ stacked_weight = weight.detach()
83
+ stacked_weight.shard_id = layer_idx
84
+ yield f"context_wkv_proj.{match.group(2)}", stacked_weight
85
86
87
class DSparkDeepseekV4Model(nn.Module):
@@ -113,6 +134,15 @@ class DSparkDeepseekV4Model(nn.Module):
134
for i in range(self.num_dspark_layers)
135
]
136
)
137
+ self.context_wkv_proj = MergedColumnParallelLinear(
138
+ config.hidden_size,
139
+ [config.head_dim] * self.num_dspark_layers,
140
+ bias=False,
141
+ return_bias=False,
142
+ quant_config=vllm_config.quant_config,
143
+ prefix=maybe_prefix(prefix, "context_wkv_proj"),
144
+ disable_tp=True,
145
+ )
146
147
# Heads: final norm + hc_head, and the Markov + confidence heads
148
# Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
@@ -173,15 +203,16 @@ class DSparkDeepseekV4Model(nn.Module):
203
place draft layers in different groups). ``None`` (or a ``None`` entry)
204
runs the projection to reserve workspace but writes nothing (profiling).
205
"""
176
- for i, layer in enumerate(self.layers):
206
+ all_kv = self.context_wkv_proj(main_x).view(
207
+ main_x.shape[0], self.num_dspark_layers, self.config.head_dim
208
+ )
209
+ for i, (layer, kv) in enumerate(
210
+ zip(self.layers, all_kv.unbind(1), strict=True)
211
+ ):
212
slot_mapping = (
213
None if context_slot_mappings is None else context_slot_mappings[i]
214
)
215
attn = layer.attn
181
- # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
182
- # (q_lora part discarded), then RoPE/quant/insert via the fused op.
183
- qr_kv, _ = attn.fused_wqa_wkv(main_x)
184
- kv = qr_kv[..., attn.q_lora_rank :]
216
kv = attn.kv_norm(kv)
217
if slot_mapping is None:
218
continue
@@ -429,6 +460,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
460
head_start = n_local_head * tp_rank
461
head_end = n_local_head * (tp_rank + 1)
462
463
+ weights = _duplicate_context_wkv_weights(weights, len(self.model.layers))
464
for name, loaded_weight in weights:
465
mapped = self._remap_dspark_name(name)
466
if mapped is None:
@@ -452,6 +484,12 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
484
self.quant_config, name, loaded_weight
485
)
486
487
+ if name.startswith("model.context_wkv_proj."):
488
+ param = params_dict[name]
489
+ param.weight_loader(param, loaded_weight, loaded_weight.shard_id)
490
+ loaded_params.add(name)
491
+ continue
492
+
493
# E8M0 expert scales: keep raw exponent bytes.
494
if ".experts." in name:
495
if (
@@ -522,6 +560,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
560
561
Returns None for non-mtp weights (owned by the target model).
562
"""
563
+ if name.startswith("context_wkv_proj."):
564
+ return f"model.{name}"
565
m = re.match(r"mtp\.(\d+)\.(.*)", name)
566
if m is None:
567
return None