[ROCm] Enable DeepSeek-V4 DSpark speculative decoding on AMD (MI350X / MI355X, gfx950) (#47419)
Signed-off-by: larryli2-amd <larryli2@amd.com> Signed-off-by: larryli2-amd <Larry.Li@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
larryli2-amd committed
Jul 10, 2026 at 23:22 UTC
c227aaa3f8edd02dae4583e27246430eebabfb25
6 files changed
+586
-14
tests/models/test_registry.py
+5
-3
@@ -48,9 +48,11 @@ def test_registry_imports(model_arch):
48
"(see #41376)"
49
)
50
51
- # DSpark draft model is NVIDIA-only; class is stubbed to None on ROCm/XPU.
52
- if model_arch == "DSparkDraftModel" and not current_platform.is_cuda():
53
- pytest.skip("DSparkDraftModel is only supported on CUDA")
51
+ # DSpark draft model is supported on CUDA and ROCm; stubbed to None on XPU.
52
+ if model_arch == "DSparkDraftModel" and not (
53
+ current_platform.is_cuda() or current_platform.is_rocm()
54
+ ):
55
+ pytest.skip("DSparkDraftModel is only supported on CUDA and ROCm")
56
57
# Ensure all model classes can be imported successfully
58
model_cls = ModelRegistry._try_load_model_cls(model_arch)
vllm/config/speculative.py
+25
@@ -942,6 +942,31 @@ class SpeculativeConfig:
942
"`num_speculative_tokens` was not provided"
943
)
944
945
+ if self.method == "dspark":
946
+ # DSpark is a semi-autoregressive *block* drafter. A
947
+ # speculative length smaller than the checkpoint's block
948
+ # feeds the block / Markov-head machinery an unsupported
949
+ # layout and yields incorrect (garbled) output rather than
950
+ # merely lower acceptance. Require num_speculative_tokens to
951
+ # be at least the block size (e.g. 5 or 7 for DeepSeek-V4).
952
+ dspark_block_size = getattr(
953
+ self.draft_model_config.hf_config,
954
+ "dspark_block_size",
955
+ None,
956
+ )
957
+ if (
958
+ dspark_block_size is not None
959
+ and self.num_speculative_tokens < dspark_block_size
960
+ ):
961
+ raise ValueError(
962
+ "DSpark requires num_speculative_tokens >= "
963
+ f"dspark_block_size ({dspark_block_size}); got "
964
+ f"{self.num_speculative_tokens}. Smaller values "
965
+ "produce incorrect output. Use "
966
+ f"num_speculative_tokens={dspark_block_size} or "
967
+ "larger (e.g. 7)."
968
+ )
969
+
970
self.draft_tensor_parallel_size = (
971
SpeculativeConfig._verify_and_get_draft_tp(
972
self.target_parallel_config,
vllm/models/deepseek_v4/__init__.py
+4
-4
@@ -15,16 +15,16 @@ from .quant_config import DeepseekV4FP8Config
15
# default that mypy sees; the ROCm/XPU branches override at runtime and are
16
# kept type-compatible via ``# type: ignore[assignment]``.
17
if current_platform.is_rocm():
18
+ from .amd.dspark import ( # type: ignore[assignment]
19
+ DSparkDeepseekV4ForCausalLM,
20
+ )
21
from .amd.model import DeepseekV4ForCausalLM
22
from .amd.mtp import DeepSeekV4MTP
20
-
21
- # DSpark is NVIDIA-only for now.
22
- DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment]
23
elif current_platform.is_xpu():
24
from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
25
from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
26
27
- DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment]
27
+ DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment, misc]
28
else:
29
from .nvidia.dspark import ( # type: ignore[assignment]
30
DSparkDeepseekV4ForCausalLM,
vllm/models/deepseek_v4/amd/dspark.py
new
+499
@@ -0,0 +1,499 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""DSpark draft model for DeepSeek-V4 on ROCm/AMD (gfx950).
4
+
5
+ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for
6
+``amd/mtp.py``:
7
+
8
+ * import ``DeepseekV4DecoderLayer`` from the AMD ``.model`` (aiter/triton
9
+ attention + MHC CustomOp path) instead of the nvidia one;
10
+ * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter /
11
+ tilelang / triton / torch) instead of calling the tilelang kernels directly,
12
+ and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter
13
+ path, where the decoder layer already applies hc_post in-layer);
14
+ * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` /
15
+ ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py).
16
+
17
+Everything else — the semi-autoregressive drafting hooks, the Markov head, the
18
+sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is
19
+pure torch / Triton and shared with the nvidia implementation unchanged.
20
+"""
21
+
22
+from collections.abc import Iterable
23
+
24
+import regex as re
25
+import torch
26
+import torch.nn as nn
27
+
28
+from vllm.config import VllmConfig, get_current_vllm_config
29
+from vllm.distributed import (
30
+ get_tensor_model_parallel_rank,
31
+ get_tensor_model_parallel_world_size,
32
+)
33
+from vllm.logger import init_logger
34
+from vllm.model_executor.layers.fused_moe import (
35
+ fused_moe_make_expert_params_mapping,
36
+)
37
+from vllm.model_executor.layers.layernorm import RMSNorm
38
+from vllm.model_executor.layers.linear import ReplicatedLinear
39
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
+from vllm.model_executor.layers.mhc import HCHeadOp
41
+from vllm.model_executor.layers.vocab_parallel_embedding import (
42
+ ParallelLMHead,
43
+ VocabParallelEmbedding,
44
+)
45
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
46
+from vllm.model_executor.models.qwen3_dspark import (
47
+ DSparkMarkovHead,
48
+)
49
+from vllm.model_executor.models.utils import maybe_prefix
50
+
51
+from .model import (
52
+ DeepseekV4DecoderLayer,
53
+)
54
+
55
+logger = init_logger(__name__)
56
+
57
+# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
58
+# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
59
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
60
+
61
+
62
+class DSparkDeepseekV4Model(nn.Module):
63
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
64
+ super().__init__()
65
+ assert vllm_config.speculative_config is not None
66
+ config = vllm_config.speculative_config.draft_model_config.hf_config
67
+ self.config = config
68
+ self.hidden_size = config.hidden_size
69
+ self.hc_mult = config.hc_mult
70
+ self.hc_eps = config.hc_eps
71
+ self.rms_norm_eps = config.rms_norm_eps
72
+ self.num_hidden_layers = config.num_hidden_layers
73
+ self.target_layer_ids = tuple(config.dspark_target_layer_ids)
74
+
75
+ self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
76
+
77
+ # Shared with the target (aliased by the speculator's loading utility).
78
+ self.embed_tokens = VocabParallelEmbedding(
79
+ config.vocab_size,
80
+ config.hidden_size,
81
+ prefix=maybe_prefix(prefix, "embed_tokens"),
82
+ )
83
+
84
+ self.main_proj = ReplicatedLinear(
85
+ config.hidden_size * len(self.target_layer_ids),
86
+ config.hidden_size,
87
+ bias=False,
88
+ return_bias=False,
89
+ quant_config=vllm_config.quant_config,
90
+ prefix=maybe_prefix(prefix, "main_proj"),
91
+ )
92
+ self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
93
+
94
+ current_vllm_config = get_current_vllm_config()
95
+ self.layers = nn.ModuleList(
96
+ [
97
+ DeepseekV4DecoderLayer(
98
+ current_vllm_config,
99
+ prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
100
+ )
101
+ for i in range(self.num_dspark_layers)
102
+ ]
103
+ )
104
+
105
+ # Heads: final norm + hc_head, and the Markov head
106
+ # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
107
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
108
+ hc_dim = self.hc_mult * config.hidden_size
109
+ self.hc_head_fn = nn.Parameter(
110
+ torch.empty(self.hc_mult, hc_dim, dtype=torch.float32),
111
+ requires_grad=False,
112
+ )
113
+ self.hc_head_base = nn.Parameter(
114
+ torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False
115
+ )
116
+ self.hc_head_scale = nn.Parameter(
117
+ torch.empty(1, dtype=torch.float32), requires_grad=False
118
+ )
119
+ draft_vocab_size = (
120
+ getattr(config, "draft_vocab_size", None) or config.vocab_size
121
+ )
122
+ self.markov_head = DSparkMarkovHead(
123
+ config.vocab_size,
124
+ draft_vocab_size,
125
+ config.dspark_markov_rank,
126
+ prefix=maybe_prefix(prefix, "markov_head"),
127
+ )
128
+
129
+ # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch),
130
+ # replacing the direct nvidia tilelang kernel call.
131
+ self.hc_head_op = HCHeadOp()
132
+
133
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
134
+ return self.embed_tokens(input_ids)
135
+
136
+ def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
137
+ """main_x = main_norm(main_proj(concat of target aux hidden states)).
138
+
139
+ ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
140
+ """
141
+ return self.main_norm(self.main_proj(aux_hidden_states))
142
+
143
+ @torch.inference_mode()
144
+ def precompute_and_store_context_kv(
145
+ self,
146
+ main_x: torch.Tensor,
147
+ context_positions: torch.Tensor,
148
+ context_slot_mappings: list[torch.Tensor | None] | None = None,
149
+ ) -> None:
150
+ """Insert the sliding-window context KV for every draft layer.
151
+
152
+ Mirrors the reference DSparkAttention: each layer derives its context KV
153
+ from the SAME projected target hidden ``main_x``, via that layer's own
154
+ ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
155
+ layer's context slots.
156
+
157
+ ``context_slot_mappings`` is a per-layer list (each entry is the context
158
+ slot mapping for that layer's kv-cache group, since the hybrid manager may
159
+ place draft layers in different groups). ``None`` (or a ``None`` entry)
160
+ runs the projection to reserve workspace but writes nothing (profiling).
161
+ """
162
+ for i, layer in enumerate(self.layers):
163
+ slot_mapping = (
164
+ None if context_slot_mappings is None else context_slot_mappings[i]
165
+ )
166
+ attn = layer.attn
167
+ # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
168
+ # (q_lora part discarded), then RoPE/quant/insert via the fused op.
169
+ qr_kv, _ = attn.fused_wqa_wkv(main_x)
170
+ kv = qr_kv[..., attn.q_lora_rank :]
171
+ kv = attn.kv_norm(kv)
172
+ if slot_mapping is None:
173
+ continue
174
+ _insert_context_kv(attn, kv, context_positions, slot_mapping)
175
+
176
+ def forward(
177
+ self,
178
+ input_ids: torch.Tensor,
179
+ positions: torch.Tensor,
180
+ inputs_embeds: torch.Tensor | None = None,
181
+ ) -> torch.Tensor:
182
+ if inputs_embeds is None:
183
+ inputs_embeds = self.embed_input_ids(input_ids)
184
+ # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
185
+ hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
186
+
187
+ residual = post_mix = res_mix = None
188
+ for layer in self.layers:
189
+ hidden_states, residual, post_mix, res_mix = layer(
190
+ hidden_states,
191
+ positions,
192
+ input_ids,
193
+ post_mix,
194
+ res_mix,
195
+ residual,
196
+ )
197
+ # On the fused-MHC path the trailing hc_post must be applied here; on the
198
+ # aiter unfused path (ROCm default) the decoder layer already applied
199
+ # hc_post in-layer and returned None mixes, so this is skipped. Mirrors
200
+ # amd/mtp.py.
201
+ last_layer = self.layers[-1]
202
+ if last_layer.use_fused_mhc:
203
+ hidden_states = last_layer.hc_post(
204
+ hidden_states, residual, post_mix, res_mix
205
+ )
206
+ # hc_head reduces the hc copies; return the PRE-norm head hidden.
207
+ hidden_states = self.hc_head_op(
208
+ hidden_states,
209
+ self.hc_head_fn,
210
+ self.hc_head_scale,
211
+ self.hc_head_base,
212
+ self.rms_norm_eps,
213
+ self.hc_eps,
214
+ )
215
+ return hidden_states
216
+
217
+
218
+def _insert_context_kv(
219
+ attn: nn.Module,
220
+ kv: torch.Tensor,
221
+ positions: torch.Tensor,
222
+ slot_mapping: torch.Tensor,
223
+) -> None:
224
+ """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
225
+
226
+ Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
227
+ query and discard it, since context tokens have no query). Mirrors
228
+ ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
229
+ """
230
+ swa_cache = attn.swa_cache_layer.kv_cache
231
+ block_size = attn.swa_cache_layer.block_size
232
+ cos_sin_cache = attn.rotary_emb.cos_sin_cache
233
+ cache_dtype = swa_cache.dtype
234
+ n_ctx = kv.shape[0]
235
+ dummy_q = torch.zeros(
236
+ (n_ctx, attn.n_local_heads, attn.head_dim),
237
+ dtype=kv.dtype,
238
+ device=kv.device,
239
+ )
240
+ if cache_dtype == torch.uint8:
241
+ # fp8_ds_mla UE8M0 paged layout (the gfx950 aiter SWA cache path).
242
+ swa_2d = swa_cache.view(swa_cache.shape[0], -1)
243
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
244
+ dummy_q,
245
+ kv,
246
+ swa_2d,
247
+ slot_mapping,
248
+ positions,
249
+ cos_sin_cache,
250
+ attn.padded_heads,
251
+ attn.eps,
252
+ block_size,
253
+ )
254
+ elif cache_dtype == torch.bfloat16:
255
+ swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
256
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
257
+ dummy_q,
258
+ kv,
259
+ swa_3d,
260
+ slot_mapping,
261
+ positions,
262
+ cos_sin_cache,
263
+ attn.eps,
264
+ block_size,
265
+ )
266
+ else: # per-tensor fp8 (torch.float8_e4m3fn)
267
+ # NOTE(rocm): unreachable on ROCm/aiter, where the SWA cache dtype is
268
+ # uint8 (fp8_ds_mla) or bfloat16. This branch relies on FlashInfer-only
269
+ # attributes (``_flashinfer_fp8_*``) that the aiter attention layer does
270
+ # not define; kept for parity with the nvidia path.
271
+ swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
272
+ dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
273
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
274
+ dummy_q,
275
+ kv,
276
+ dummy_q_fp8,
277
+ swa_3d,
278
+ slot_mapping,
279
+ positions,
280
+ cos_sin_cache,
281
+ attn._flashinfer_fp8_kv_scale,
282
+ attn._flashinfer_fp8_q_scale_inv,
283
+ attn.eps,
284
+ block_size,
285
+ )
286
+
287
+
288
+class DSparkDeepseekV4ForCausalLM(nn.Module):
289
+ # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
290
+ # load_dspark_model always aliases the target's.
291
+ has_own_embed_tokens = False
292
+ has_own_lm_head = False
293
+
294
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
295
+ super().__init__()
296
+ assert vllm_config.speculative_config is not None
297
+ self.draft_model_config = vllm_config.speculative_config.draft_model_config
298
+ self.config = self.draft_model_config.hf_config
299
+ self.model = DSparkDeepseekV4Model(
300
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
301
+ )
302
+ # Shared with the target (aliased by the speculator's load utility).
303
+ self.lm_head = ParallelLMHead(
304
+ self.config.vocab_size,
305
+ self.config.hidden_size,
306
+ prefix=maybe_prefix(prefix, "lm_head"),
307
+ )
308
+ self.logits_processor = LogitsProcessor(self.config.vocab_size)
309
+
310
+ # --- Hooks used by the speculator -------------------------------------
311
+
312
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
313
+ return self.model.embed_input_ids(input_ids)
314
+
315
+ def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
316
+ return self.model.combine_hidden_states(aux_hidden_states)
317
+
318
+ def get_draft_kv_cache_layer_names(self) -> list[str]:
319
+ # DSV4 MLA path: each draft layer's sliding-window cache is a separate
320
+ # layer, named by its prefix.
321
+ return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
322
+
323
+ def precompute_and_store_context_kv(
324
+ self,
325
+ context_states: torch.Tensor,
326
+ context_positions: torch.Tensor,
327
+ context_slot_mappings: list[torch.Tensor | None] | None = None,
328
+ ) -> None:
329
+ self.model.precompute_and_store_context_kv(
330
+ context_states, context_positions, context_slot_mappings
331
+ )
332
+
333
+ def forward(
334
+ self,
335
+ input_ids: torch.Tensor,
336
+ positions: torch.Tensor,
337
+ inputs_embeds: torch.Tensor | None = None,
338
+ ) -> torch.Tensor:
339
+ # Returns the pre-norm hc_head hidden ([T, hidden_size]).
340
+ return self.model(input_ids, positions, inputs_embeds)
341
+
342
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
343
+ """Base logits U_k = lm_head(norm(head_hidden))."""
344
+ return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
345
+
346
+ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
347
+ # Full-vocab draft: base logits, no d2t scatter.
348
+ return self.compute_logits(hidden_states)
349
+
350
+ def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
351
+ return draft_ids # full-vocab: draft ids are target ids
352
+
353
+ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
354
+ return self.model.markov_head.embed(token_ids)
355
+
356
+ def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
357
+ return self.model.markov_head.bias(markov_embed, self.logits_processor)
358
+
359
+ # --- Weight loading ----------------------------------------------------
360
+
361
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
362
+ """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
363
+
364
+ Non-mtp weights (embed/head/main layers) belong to the target model and
365
+ are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
366
+ """
367
+ # AMD DeepseekV4MoE has no mega-MoE path; always use the standard
368
+ # per-expert fused-MoE mapping (mirrors amd/mtp.py).
369
+ expert_mapping = fused_moe_make_expert_params_mapping(
370
+ self,
371
+ ckpt_gate_proj_name="w1",
372
+ ckpt_down_proj_name="w2",
373
+ ckpt_up_proj_name="w3",
374
+ num_experts=self.config.n_routed_experts,
375
+ )
376
+ expert_scale_suffix = (
377
+ ".weight_scale"
378
+ if getattr(self.config, "expert_dtype", "fp4") == "fp4"
379
+ else ".weight_scale_inv"
380
+ )
381
+
382
+ # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
383
+ stacked_params_mapping = [
384
+ ("gate_up_proj", "w1", 0),
385
+ ("gate_up_proj", "w3", 1),
386
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
387
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
388
+ ]
389
+
390
+ params_dict = dict(self.named_parameters())
391
+ loaded_params: set[str] = set()
392
+
393
+ tp_size = get_tensor_model_parallel_world_size()
394
+ tp_rank = get_tensor_model_parallel_rank()
395
+ n_local_head = self.config.num_attention_heads // tp_size
396
+ head_start = n_local_head * tp_rank
397
+ head_end = n_local_head * (tp_rank + 1)
398
+
399
+ for name, loaded_weight in weights:
400
+ mapped = self._remap_dspark_name(name)
401
+ if mapped is None:
402
+ continue
403
+ name = mapped
404
+
405
+ # ``.scale`` -> per-method scale suffix.
406
+ if name.endswith(".scale"):
407
+ suffix = (
408
+ expert_scale_suffix
409
+ if _EXPERT_SCALE_RE.search(name)
410
+ else ".weight_scale_inv"
411
+ )
412
+ name = name.removesuffix(".scale") + suffix
413
+
414
+ # E8M0 expert scales: keep raw exponent bytes.
415
+ if ".experts." in name:
416
+ if (
417
+ "weight_scale" in name
418
+ and loaded_weight.dtype == torch.float8_e8m0fnu
419
+ ):
420
+ loaded_weight = loaded_weight.view(torch.uint8)
421
+ for param_name, weight_name, expert_id, shard_id in expert_mapping:
422
+ if weight_name not in name:
423
+ continue
424
+ name_mapped = name.replace(weight_name, param_name)
425
+ param = params_dict[name_mapped]
426
+ success = param.weight_loader(
427
+ param,
428
+ loaded_weight,
429
+ name_mapped,
430
+ shard_id=shard_id,
431
+ expert_id=expert_id,
432
+ return_success=True,
433
+ )
434
+ if success:
435
+ loaded_params.add(name_mapped)
436
+ break
437
+ continue
438
+
439
+ # Stacked rules only apply to decoder-layer weights. Head-stack params
440
+ # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
441
+ # "markov_w1" would collide with the "w1" shard rule.
442
+ is_layer_param = name.startswith("model.layers.")
443
+ for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
444
+ if not is_layer_param or weight_name not in name:
445
+ continue
446
+ name = name.replace(weight_name, param_name)
447
+ param = params_dict[name]
448
+ param.weight_loader(param, loaded_weight, stacked_shard_id)
449
+ loaded_params.add(name)
450
+ break
451
+ else:
452
+ if "attn_sink" in name:
453
+ narrow = loaded_weight[head_start:head_end]
454
+ params_dict[name][: narrow.shape[0]].copy_(narrow)
455
+ loaded_params.add(name)
456
+ continue
457
+ if ".shared_experts.w2" in name:
458
+ name = name.replace(
459
+ ".shared_experts.w2", ".shared_experts.down_proj"
460
+ )
461
+ if name.endswith(".ffn.gate.bias"):
462
+ name = name.replace(
463
+ ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
464
+ )
465
+ param = params_dict[name]
466
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
467
+ weight_loader(param, loaded_weight)
468
+ loaded_params.add(name)
469
+
470
+ logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
471
+ return loaded_params
472
+
473
+ def _remap_dspark_name(self, name: str) -> str | None:
474
+ """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
475
+
476
+ Returns None for non-mtp weights (owned by the target model).
477
+ """
478
+ m = re.match(r"mtp\.(\d+)\.(.*)", name)
479
+ if m is None:
480
+ return None
481
+ stage = int(m.group(1))
482
+ rest = m.group(2)
483
+ # The confidence head is not wired into inference yet; drop its weights.
484
+ if rest.startswith("confidence_head."):
485
+ return None
486
+ # Head-stack params live at model level (mtp.last), context combiner at
487
+ # model level (mtp.0); everything else is a per-layer decoder block.
488
+ head_prefixes = (
489
+ "norm.",
490
+ "hc_head_fn",
491
+ "hc_head_base",
492
+ "hc_head_scale",
493
+ "markov_head.",
494
+ )
495
+ if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
496
+ head_prefixes
497
+ ):
498
+ return f"model.{rest}"
499
+ return f"model.layers.{stage}.{rest}"
vllm/models/deepseek_v4/amd/model.py
+45
-5
@@ -40,7 +40,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
40
VocabParallelEmbedding,
41
)
42
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43
-from vllm.model_executor.models.interfaces import SupportsPP
43
+from vllm.model_executor.models.interfaces import (
44
+ EagleModelMixin,
45
+ SupportsEagle3,
46
+ SupportsPP,
47
+)
48
from vllm.model_executor.models.utils import (
49
AutoWeightsLoader,
50
PPMissingLayer,
@@ -437,7 +441,7 @@ class DeepseekV4DecoderLayer(nn.Module):
441
)
442
443
440
-class DeepseekV4Model(nn.Module):
444
+class DeepseekV4Model(nn.Module, EagleModelMixin):
445
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
446
super().__init__()
447
@@ -573,7 +577,19 @@ class DeepseekV4Model(nn.Module):
577
hidden_states = intermediate_tensors["hidden_states"]
578
579
residual, post_mix, res_mix = None, None, None
576
- for layer in islice(self.layers, self.start_layer, self.end_layer):
580
+ # EAGLE3 / DSpark / DFlash aux hidden states: reconstructed (post-mhc)
581
+ # hidden state at the configured target layers, averaged over the
582
+ # hc_mult streams to [T, hidden_size]. Empty unless a draft model set
583
+ # aux_hidden_state_layers.
584
+ aux_hidden_states: list[torch.Tensor] = []
585
+ # On the fused path the final layer's hc_post output is reused below
586
+ # (avoids computing hc_post twice when the last layer is also an aux
587
+ # layer).
588
+ final_aux_recon: torch.Tensor | None = None
589
+ for idx, layer in enumerate(
590
+ islice(self.layers, self.start_layer, self.end_layer),
591
+ start=self.start_layer,
592
+ ):
593
hidden_states, residual, post_mix, res_mix = layer(
594
hidden_states,
595
positions,
@@ -582,8 +598,30 @@ class DeepseekV4Model(nn.Module):
598
res_mix,
599
residual,
600
)
601
+ if (idx + 1) in self.aux_hidden_state_layers:
602
+ # On the unfused (aiter) path the layer already applied hc_post,
603
+ # so hidden_states is the reconstructed stream; on the fused
604
+ # path reconstruct it via hc_post before averaging.
605
+ if layer.use_fused_mhc:
606
+ aux_recon = layer.hc_post(
607
+ hidden_states, residual, post_mix, res_mix
608
+ )
609
+ final_aux_recon = aux_recon
610
+ else:
611
+ aux_recon = hidden_states
612
+ aux_hidden_states.append(aux_recon.mean(dim=1))
613
if layer is not None and layer.use_fused_mhc:
586
- hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
614
+ # Reuse the last layer's hc_post output if it was already computed
615
+ # for the aux hidden state above; otherwise compute it now.
616
+ if (
617
+ final_aux_recon is not None
618
+ and self.end_layer in self.aux_hidden_state_layers
619
+ ):
620
+ hidden_states = final_aux_recon
621
+ else:
622
+ hidden_states = layer.hc_post(
623
+ hidden_states, residual, post_mix, res_mix
624
+ )
625
626
if not get_pp_group().is_last_rank:
627
return IntermediateTensors({"hidden_states": hidden_states})
@@ -601,6 +639,8 @@ class DeepseekV4Model(nn.Module):
639
self.hc_eps,
640
)
641
hidden_states = self.norm(hidden_states)
642
+ if len(aux_hidden_states) > 0:
643
+ return hidden_states, aux_hidden_states
644
return hidden_states
645
646
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -751,7 +791,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
791
)
792
793
754
-class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
794
+class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
795
model_cls = DeepseekV4Model
796
797
# Default mapper assumes the original FP4-expert checkpoint layout.
vllm/models/deepseek_v4/amd/rocm.py
+8
-2
@@ -371,8 +371,12 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
371
def __init__(self, *args, **kwargs):
372
super().__init__(*args, **kwargs)
373
max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens
374
+ # The non-causal (DSpark draft) path widens each token's SWA index list
375
+ # to ``noncausal_index_width`` (>= window_size), so size the persistent
376
+ # ragged buffer to the wider bound to cover both causal and non-causal.
377
+ swa_index_width = max(self.window_size, self.noncausal_index_width)
378
self.decode_swa_ragged_indices_buffer = torch.empty(
375
- max_tokens * self.window_size,
379
+ max_tokens * swa_index_width,
380
dtype=torch.int32,
381
device=self.device,
382
)
@@ -411,7 +415,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
415
self.decode_swa_ragged_indices_buffer,
416
self.decode_swa_ragged_indptr_buffer,
417
base.num_decode_tokens,
414
- self.window_size,
418
+ # Actual dense width for this build: window_size (causal) or
419
+ # noncausal_index_width (DSpark non-causal draft).
420
+ base.decode_swa_indices.shape[-1],
421
)
422
423
return DeepseekV4ROCMAiterSparseSWAMetadata(