[XPU] Add DSpark speculative decoding support for DeepSeek-V4 (#47677)
Signed-off-by: Ma Jian <jian1.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ma Jian committed
Jul 16, 2026 at 08:59 UTC
9d1c695be5875c49e34b63d7b1f9aad282551c30
3 files changed
+454
-6
vllm/models/deepseek_v4/__init__.py
+1
-2
@@ -21,10 +21,9 @@ if current_platform.is_rocm():
21
from .amd.model import DeepseekV4ForCausalLM
22
from .amd.mtp import DeepSeekV4MTP
23
elif current_platform.is_xpu():
24
+ from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment]
25
from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
26
from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
26
-
27
- DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment, misc]
27
else:
28
from .nvidia.dspark import ( # type: ignore[assignment]
29
DSparkDeepseekV4ForCausalLM,
vllm/models/deepseek_v4/xpu/dspark.py
new
+436
@@ -0,0 +1,436 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""DSpark draft model for DeepSeek-V4 on Intel XPU.
4
+
5
+Minimal XPU port of nvidia/dspark.py. Replaces tilelang MHC kernels with
6
+the platform-agnostic custom ops (HCHeadOp, MHCPostOp) already used by the
7
+XPU MTP path, and uses the XPU Triton-based qnorm_rope_kv_fp8_insert for
8
+context KV precomputation.
9
+"""
10
+
11
+from collections.abc import Iterable
12
+
13
+import regex as re
14
+import torch
15
+import torch.nn as nn
16
+
17
+from vllm.config import VllmConfig, get_current_vllm_config
18
+from vllm.distributed import (
19
+ get_tensor_model_parallel_rank,
20
+ get_tensor_model_parallel_world_size,
21
+)
22
+from vllm.logger import init_logger
23
+from vllm.model_executor.layers.fused_moe import (
24
+ fused_moe_make_expert_params_mapping,
25
+)
26
+from vllm.model_executor.layers.layernorm import RMSNorm
27
+from vllm.model_executor.layers.linear import ReplicatedLinear
28
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
29
+from vllm.model_executor.layers.mhc import HCHeadOp, MHCPostOp
30
+from vllm.model_executor.layers.vocab_parallel_embedding import (
31
+ ParallelLMHead,
32
+ VocabParallelEmbedding,
33
+)
34
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35
+from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
36
+from vllm.model_executor.models.utils import maybe_prefix
37
+
38
+from .model import (
39
+ DeepseekV4DecoderLayer,
40
+ make_deepseek_v4_expert_params_mapping,
41
+)
42
+
43
+logger = init_logger(__name__)
44
+
45
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
46
+
47
+
48
+class DSparkDeepseekV4Model(nn.Module):
49
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
50
+ super().__init__()
51
+ assert vllm_config.speculative_config is not None
52
+ config = vllm_config.speculative_config.draft_model_config.hf_config
53
+ self.config = config
54
+ self.hidden_size = config.hidden_size
55
+ self.hc_mult = config.hc_mult
56
+ self.hc_eps = config.hc_eps
57
+ self.rms_norm_eps = config.rms_norm_eps
58
+ self.num_hidden_layers = config.num_hidden_layers
59
+ self.target_layer_ids = tuple(config.dspark_target_layer_ids)
60
+
61
+ self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
62
+
63
+ # Shared with target (aliased by speculator loading utility).
64
+ self.embed_tokens = VocabParallelEmbedding(
65
+ config.vocab_size,
66
+ config.hidden_size,
67
+ prefix=maybe_prefix(prefix, "embed_tokens"),
68
+ )
69
+
70
+ self.main_proj = ReplicatedLinear(
71
+ config.hidden_size * len(self.target_layer_ids),
72
+ config.hidden_size,
73
+ bias=False,
74
+ return_bias=False,
75
+ quant_config=vllm_config.quant_config,
76
+ prefix=maybe_prefix(prefix, "main_proj"),
77
+ )
78
+ self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
79
+
80
+ current_vllm_config = get_current_vllm_config()
81
+ self.layers = nn.ModuleList(
82
+ [
83
+ DeepseekV4DecoderLayer(
84
+ current_vllm_config,
85
+ prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
86
+ )
87
+ for i in range(self.num_dspark_layers)
88
+ ]
89
+ )
90
+
91
+ # Heads: final norm + hc_head, and the Markov head
92
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
93
+ hc_dim = self.hc_mult * config.hidden_size
94
+ self.hc_head_fn = nn.Parameter(
95
+ torch.empty(self.hc_mult, hc_dim, dtype=torch.float32),
96
+ requires_grad=False,
97
+ )
98
+ self.hc_head_base = nn.Parameter(
99
+ torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False
100
+ )
101
+ self.hc_head_scale = nn.Parameter(
102
+ torch.empty(1, dtype=torch.float32), requires_grad=False
103
+ )
104
+ draft_vocab_size = (
105
+ getattr(config, "draft_vocab_size", None) or config.vocab_size
106
+ )
107
+ self.markov_head = DSparkMarkovHead(
108
+ config.vocab_size,
109
+ draft_vocab_size,
110
+ config.dspark_markov_rank,
111
+ prefix=maybe_prefix(prefix, "markov_head"),
112
+ )
113
+
114
+ # XPU MHC ops (replaces tilelang)
115
+ self.mhc_post_op = MHCPostOp()
116
+ self.hc_head_op = HCHeadOp()
117
+
118
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
119
+ return self.embed_tokens(input_ids)
120
+
121
+ def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
122
+ """main_x = main_norm(main_proj(concat of target aux hidden states))."""
123
+ return self.main_norm(self.main_proj(aux_hidden_states))
124
+
125
+ @torch.inference_mode()
126
+ def precompute_and_store_context_kv(
127
+ self,
128
+ main_x: torch.Tensor,
129
+ context_positions: torch.Tensor,
130
+ context_slot_mappings: list[torch.Tensor | None] | None = None,
131
+ ) -> None:
132
+ """Insert the sliding-window context KV for every draft layer.
133
+
134
+ Each layer derives its context KV from the SAME projected target hidden
135
+ ``main_x``, via that layer's own wkv + kv_norm + RoPE + quant, then
136
+ writes it at the layer's context slots.
137
+ """
138
+ for i, layer in enumerate(self.layers):
139
+ slot_mapping = (
140
+ None if context_slot_mappings is None else context_slot_mappings[i]
141
+ )
142
+ attn = layer.attn
143
+ # wkv part of the fused wq_a|wkv projection (q_lora part discarded)
144
+ qr_kv, _ = attn.fused_wqa_wkv(main_x)
145
+ kv = qr_kv[..., attn.q_lora_rank :]
146
+ kv = attn.kv_norm(kv)
147
+ if slot_mapping is None:
148
+ continue
149
+ _insert_context_kv(attn, kv, context_positions, slot_mapping)
150
+
151
+ def forward(
152
+ self,
153
+ input_ids: torch.Tensor,
154
+ positions: torch.Tensor,
155
+ inputs_embeds: torch.Tensor | None = None,
156
+ ) -> torch.Tensor:
157
+ if inputs_embeds is None:
158
+ inputs_embeds = self.embed_input_ids(input_ids)
159
+ # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
160
+ hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
161
+
162
+ residual = post_mix = res_mix = None
163
+ for layer in self.layers:
164
+ hidden_states, residual, post_mix, res_mix = layer(
165
+ hidden_states,
166
+ positions,
167
+ input_ids,
168
+ post_mix,
169
+ res_mix,
170
+ residual,
171
+ )
172
+ # mhc_post: merge hyper-connection copies
173
+ hidden_states = self.mhc_post_op(hidden_states, residual, post_mix, res_mix)
174
+ # hc_head: reduces hc copies; return pre-norm head hidden
175
+ hidden_states = self.hc_head_op(
176
+ hidden_states,
177
+ self.hc_head_fn,
178
+ self.hc_head_scale,
179
+ self.hc_head_base,
180
+ self.rms_norm_eps,
181
+ self.hc_eps,
182
+ )
183
+ return hidden_states
184
+
185
+
186
+def _insert_context_kv(
187
+ attn: nn.Module,
188
+ kv: torch.Tensor,
189
+ positions: torch.Tensor,
190
+ slot_mapping: torch.Tensor,
191
+) -> None:
192
+ """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
193
+
194
+ On XPU, we reuse the same xpu_qnorm_rope_kv_fp8_insert kernel used in
195
+ the forward path, passing a dummy q (result discarded).
196
+ """
197
+ from vllm.models.deepseek_v4.xpu.xpu_qnorm_rope_kv_fp8_insert import (
198
+ xpu_qnorm_rope_kv_fp8_insert,
199
+ )
200
+
201
+ swa_cache = attn.swa_cache_layer.kv_cache
202
+ block_size = attn.swa_cache_layer.block_size
203
+ cos_sin_cache = attn.rotary_emb.cos_sin_cache
204
+ n_ctx = kv.shape[0]
205
+
206
+ # Dummy q — we only care about the KV insert side effect.
207
+ dummy_q = torch.empty(
208
+ (n_ctx, attn.n_local_heads, attn.head_dim),
209
+ dtype=kv.dtype,
210
+ device=kv.device,
211
+ )
212
+ xpu_qnorm_rope_kv_fp8_insert(
213
+ dummy_q,
214
+ kv,
215
+ swa_cache,
216
+ slot_mapping,
217
+ positions,
218
+ cos_sin_cache,
219
+ attn.eps,
220
+ block_size,
221
+ )
222
+
223
+
224
+class DSparkDeepseekV4ForCausalLM(nn.Module):
225
+ """XPU DSpark draft model entry point for DeepSeek-V4."""
226
+
227
+ has_own_embed_tokens = False
228
+ has_own_lm_head = False
229
+
230
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
231
+ super().__init__()
232
+ assert vllm_config.speculative_config is not None
233
+ self.draft_model_config = vllm_config.speculative_config.draft_model_config
234
+ self.config = self.draft_model_config.hf_config
235
+ self.model = DSparkDeepseekV4Model(
236
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
237
+ )
238
+ # Shared with the target (aliased by the speculator's load utility).
239
+ self.lm_head = ParallelLMHead(
240
+ self.config.vocab_size,
241
+ self.config.hidden_size,
242
+ prefix=maybe_prefix(prefix, "lm_head"),
243
+ )
244
+ self.logits_processor = LogitsProcessor(self.config.vocab_size)
245
+
246
+ # --- Hooks used by the speculator ---
247
+
248
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
249
+ return self.model.embed_input_ids(input_ids)
250
+
251
+ def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
252
+ return self.model.combine_hidden_states(aux_hidden_states)
253
+
254
+ def get_draft_kv_cache_layer_names(self) -> list[str]:
255
+ return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
256
+
257
+ def precompute_and_store_context_kv(
258
+ self,
259
+ context_states: torch.Tensor,
260
+ context_positions: torch.Tensor,
261
+ context_slot_mappings: list[torch.Tensor | None] | None = None,
262
+ ) -> None:
263
+ self.model.precompute_and_store_context_kv(
264
+ context_states, context_positions, context_slot_mappings
265
+ )
266
+
267
+ def forward(
268
+ self,
269
+ input_ids: torch.Tensor,
270
+ positions: torch.Tensor,
271
+ inputs_embeds: torch.Tensor | None = None,
272
+ ) -> torch.Tensor:
273
+ return self.model(input_ids, positions, inputs_embeds)
274
+
275
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
276
+ """Base logits U_k = lm_head(norm(head_hidden))."""
277
+ return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
278
+
279
+ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
280
+ return self.compute_logits(hidden_states)
281
+
282
+ def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
283
+ return draft_ids # full-vocab: draft ids are target ids
284
+
285
+ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
286
+ return self.model.markov_head.embed(token_ids)
287
+
288
+ def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
289
+ return self.model.markov_head.bias(markov_embed, self.logits_processor)
290
+
291
+ # --- Weight loading ---
292
+
293
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
294
+ """Load ``mtp.{0,1,2}.*`` draft weights from the target checkpoint."""
295
+ first_layer = self.model.layers[0]
296
+ use_mega_moe = first_layer.ffn.use_mega_moe
297
+ if use_mega_moe:
298
+ expert_mapping = make_deepseek_v4_expert_params_mapping(
299
+ self.config.n_routed_experts
300
+ )
301
+ else:
302
+ expert_mapping = fused_moe_make_expert_params_mapping(
303
+ self,
304
+ ckpt_gate_proj_name="w1",
305
+ ckpt_down_proj_name="w2",
306
+ ckpt_up_proj_name="w3",
307
+ num_experts=self.config.n_routed_experts,
308
+ )
309
+ expert_scale_suffix = (
310
+ ".weight_scale"
311
+ if getattr(self.config, "expert_dtype", "fp4") == "fp4"
312
+ else ".weight_scale_inv"
313
+ )
314
+
315
+ stacked_params_mapping = [
316
+ ("gate_up_proj", "w1", 0),
317
+ ("gate_up_proj", "w3", 1),
318
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
319
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
320
+ ]
321
+
322
+ params_dict = dict(self.named_parameters())
323
+ loaded_params: set[str] = set()
324
+
325
+ tp_size = get_tensor_model_parallel_world_size()
326
+ tp_rank = get_tensor_model_parallel_rank()
327
+ n_local_head = self.config.num_attention_heads // tp_size
328
+ head_start = n_local_head * tp_rank
329
+ head_end = n_local_head * (tp_rank + 1)
330
+
331
+ for name, loaded_weight in weights:
332
+ mapped = self._remap_dspark_name(name)
333
+ if mapped is None:
334
+ continue
335
+ name = mapped
336
+
337
+ # .scale -> per-method scale suffix
338
+ if name.endswith(".scale"):
339
+ suffix = (
340
+ expert_scale_suffix
341
+ if _EXPERT_SCALE_RE.search(name)
342
+ else ".weight_scale_inv"
343
+ )
344
+ name = name.removesuffix(".scale") + suffix
345
+
346
+ # Expert weights
347
+ if ".experts." in name:
348
+ if (
349
+ "weight_scale" in name
350
+ and loaded_weight.dtype == torch.float8_e8m0fnu
351
+ ):
352
+ loaded_weight = loaded_weight.view(torch.uint8)
353
+ for param_name, weight_name, expert_id, shard_id in expert_mapping:
354
+ if weight_name not in name:
355
+ continue
356
+ name_mapped = name.replace(weight_name, param_name)
357
+ if name_mapped not in params_dict:
358
+ continue
359
+ param = params_dict[name_mapped]
360
+ success = param.weight_loader(
361
+ param,
362
+ loaded_weight,
363
+ name_mapped,
364
+ shard_id=shard_id,
365
+ expert_id=expert_id,
366
+ return_success=True,
367
+ )
368
+ if success:
369
+ loaded_params.add(name_mapped)
370
+ break
371
+ continue
372
+
373
+ # Stacked params (decoder-layer only)
374
+ is_layer_param = name.startswith("model.layers.")
375
+ for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
376
+ if not is_layer_param or weight_name not in name:
377
+ continue
378
+ name = name.replace(weight_name, param_name)
379
+ if name not in params_dict:
380
+ break
381
+ param = params_dict[name]
382
+ param.weight_loader(param, loaded_weight, stacked_shard_id)
383
+ loaded_params.add(name)
384
+ break
385
+ else:
386
+ if "attn_sink" in name:
387
+ if name not in params_dict:
388
+ continue
389
+ narrow = loaded_weight[head_start:head_end]
390
+ params_dict[name][: narrow.shape[0]].copy_(narrow)
391
+ loaded_params.add(name)
392
+ continue
393
+ if ".shared_experts.w2" in name:
394
+ name = name.replace(
395
+ ".shared_experts.w2", ".shared_experts.down_proj"
396
+ )
397
+ if name.endswith(".ffn.gate.bias"):
398
+ name = name.replace(
399
+ ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
400
+ )
401
+ if name not in params_dict:
402
+ continue
403
+ param = params_dict[name]
404
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
405
+ weight_loader(param, loaded_weight)
406
+ loaded_params.add(name)
407
+
408
+ self._finalize_moe()
409
+ logger.info_once("DSpark XPU draft model loaded: %d params", len(loaded_params))
410
+ return loaded_params
411
+
412
+ def _finalize_moe(self) -> None:
413
+ for layer in self.model.layers:
414
+ layer.ffn.finalize_mega_moe_weights()
415
+
416
+ def _remap_dspark_name(self, name: str) -> str | None:
417
+ """Map checkpoint ``mtp.{i}.*`` name to this model's parameter path."""
418
+ m = re.match(r"mtp\.(\d+)\.(.*)", name)
419
+ if m is None:
420
+ return None
421
+ stage = int(m.group(1))
422
+ rest = m.group(2)
423
+ if rest.startswith("confidence_head."):
424
+ return None
425
+ head_prefixes = (
426
+ "norm.",
427
+ "hc_head_fn",
428
+ "hc_head_base",
429
+ "hc_head_scale",
430
+ "markov_head.",
431
+ )
432
+ if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
433
+ head_prefixes
434
+ ):
435
+ return f"model.{rest}"
436
+ return f"model.layers.{stage}.{rest}"
vllm/models/deepseek_v4/xpu/model.py
+17
-4
@@ -44,7 +44,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
44
VocabParallelEmbedding,
45
)
46
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
47
-from vllm.model_executor.models.interfaces import SupportsPP
47
+from vllm.model_executor.models.interfaces import (
48
+ EagleModelMixin,
49
+ SupportsEagle3,
50
+ SupportsPP,
51
+)
52
from vllm.model_executor.models.utils import (
53
AutoWeightsLoader,
54
PPMissingLayer,
@@ -975,7 +979,7 @@ class DeepseekV4DecoderLayer(nn.Module):
979
980
981
@support_torch_compile
978
-class DeepseekV4Model(nn.Module):
982
+class DeepseekV4Model(nn.Module, EagleModelMixin):
983
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
984
super().__init__()
985
@@ -1113,7 +1117,11 @@ class DeepseekV4Model(nn.Module):
1117
input_ids = input_ids.to(torch.int64)
1118
1119
residual, post_mix, res_mix = None, None, None
1116
- for layer in islice(self.layers, self.start_layer, self.end_layer):
1120
+ aux_hidden_states: list[torch.Tensor] = []
1121
+ for idx, layer in enumerate(
1122
+ islice(self.layers, self.start_layer, self.end_layer),
1123
+ start=self.start_layer,
1124
+ ):
1125
hidden_states, residual, post_mix, res_mix = layer(
1126
hidden_states,
1127
positions,
@@ -1122,6 +1130,9 @@ class DeepseekV4Model(nn.Module):
1130
res_mix,
1131
residual,
1132
)
1133
+ if idx + 1 in self.aux_hidden_state_layers:
1134
+ aux_recon = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1135
+ aux_hidden_states.append(aux_recon.mean(dim=1))
1136
# The fused path defers the final hc_post to the next layer's
1137
# fused_post_pre. After the last layer we must apply it explicitly.
1138
if layer is not None:
@@ -1143,6 +1154,8 @@ class DeepseekV4Model(nn.Module):
1154
self.hc_eps,
1155
)
1156
hidden_states = self.norm(hidden_states)
1157
+ if len(aux_hidden_states) > 0:
1158
+ return hidden_states, aux_hidden_states
1159
return hidden_states
1160
1161
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1300,7 +1313,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1313
)
1314
1315
1303
-class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
1316
+class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
1317
model_cls = DeepseekV4Model
1318
1319
# Default mapper assumes the original FP4-expert checkpoint layout.