@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 65e655d295

[GLM-5] Add DSV3.2/GLM5 to `vllm/models/` (#46808)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed Jun 26, 2026 at 14:09 UTC 65e655d2959111d508ad97515c85be0627a7b916
5 files changed +1170
vllm/models/deepseek_v32/__init__.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek V3.2 (``deepseek_v32``) model — hardware-isolated entry point.
4 +
5 +DeepSeek V3.2 introduced the DeepSeek Sparse Attention (DSA) architecture:
6 +MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA
7 +attend. The same model code serves any DSA checkpoint, including GLM-5.2
8 +(``glm_moe_dsa``), which reuses this architecture.
9 +"""
10 +
11 +from vllm.platforms import current_platform
12 +
13 +if current_platform.is_rocm() or current_platform.is_xpu():
14 + raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.")
15 +
16 +from .nvidia.model import DeepseekV32ForCausalLM
17 +from .nvidia.mtp import DeepseekV32MTP
18 +
19 +__all__ = [
20 + "DeepseekV32ForCausalLM",
21 + "DeepseekV32MTP",
22 +]
vllm/models/deepseek_v32/nvidia/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v32/nvidia/attention.py new
+423
@@ -0,0 +1,423 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import TYPE_CHECKING
4 +
5 +import torch
6 +import torch.nn as nn
7 +from transformers import DeepseekV2Config, DeepseekV3Config
8 +
9 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
10 +from vllm.config import CacheConfig, VllmConfig
11 +from vllm.distributed import get_tensor_model_parallel_world_size
12 +from vllm.forward_context import get_forward_context
13 +from vllm.model_executor.layers.attention import MLAAttention
14 +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
15 +from vllm.model_executor.layers.linear import (
16 + ColumnParallelLinear,
17 + MergedColumnParallelLinear,
18 + ReplicatedLinear,
19 + RowParallelLinear,
20 +)
21 +from vllm.model_executor.layers.quantization import QuantizationConfig
22 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
23 + per_token_group_quant_fp8,
24 +)
25 +from vllm.model_executor.layers.rotary_embedding import get_rope
26 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
27 +from vllm.model_executor.models.deepseek_v2 import (
28 + DeepSeekV2FusedQkvAProjLinear,
29 + DeepseekV32IndexerCache,
30 + yarn_get_mscale,
31 +)
32 +from vllm.model_executor.models.utils import extract_layer_index
33 +from vllm.utils.torch_utils import is_quantized_kv_cache
34 +
35 +if TYPE_CHECKING:
36 + from vllm.model_executor.layers.attention.mla_attention import (
37 + MLACommonMetadata,
38 + )
39 +
40 +
41 +class DeepseekV32Indexer(nn.Module):
42 + def __init__(
43 + self,
44 + vllm_config: VllmConfig,
45 + config: DeepseekV2Config | DeepseekV3Config,
46 + hidden_size: int,
47 + q_lora_rank: int,
48 + quant_config: QuantizationConfig | None,
49 + cache_config: CacheConfig | None,
50 + topk_indices_buffer: torch.Tensor | None,
51 + prefix: str = "",
52 + ):
53 + super().__init__()
54 + self.topk_tokens = config.index_topk
55 + self.n_head = config.index_n_heads
56 + self.head_dim = config.index_head_dim
57 + self.rope_dim = config.qk_rope_head_dim
58 + self.q_lora_rank = q_lora_rank
59 +
60 + # No tensor parallel, just replicated.
61 + self.wq_b = ReplicatedLinear(
62 + self.q_lora_rank,
63 + self.head_dim * self.n_head,
64 + bias=False,
65 + quant_config=quant_config,
66 + prefix=f"{prefix}.wq_b",
67 + )
68 + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
69 + # FP8 wk weights are upcasted to BF16 during loading to keep this fused.
70 + self.wk_weights_proj = MergedColumnParallelLinear(
71 + hidden_size,
72 + [self.head_dim, self.n_head],
73 + bias=False,
74 + quant_config=None,
75 + disable_tp=True,
76 + prefix=f"{prefix}.wk_weights_proj",
77 + )
78 + self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
79 + self.softmax_scale = self.head_dim**-0.5
80 +
81 + self.scale_fmt = "ue8m0"
82 + self.quant_block_size = 128
83 + self.topk_indices_buffer = topk_indices_buffer
84 +
85 + # fp8 naive cache: value in fp8 + fp32 scale per quant_block_size element.
86 + assert cache_config is not None, "DeepSeek V3.2 indexer requires cache_config"
87 + self.k_cache = DeepseekV32IndexerCache(
88 + head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
89 + dtype=torch.uint8,
90 + prefix=f"{prefix}.k_cache",
91 + cache_config=cache_config,
92 + )
93 + self.max_model_len = vllm_config.model_config.max_model_len
94 + self.prefix = prefix
95 +
96 + from vllm.v1.attention.backends.mla.indexer import (
97 + get_max_prefill_buffer_size,
98 + )
99 +
100 + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
101 + self.indexer_op = SparseAttnIndexer(
102 + self.k_cache,
103 + self.quant_block_size,
104 + self.scale_fmt,
105 + self.topk_tokens,
106 + self.head_dim,
107 + self.max_model_len,
108 + self.max_total_seq_len,
109 + self.topk_indices_buffer,
110 + )
111 +
112 + def forward(
113 + self,
114 + hidden_states: torch.Tensor,
115 + qr: torch.Tensor,
116 + positions: torch.Tensor,
117 + rotary_emb: nn.Module,
118 + ) -> torch.Tensor:
119 + q, _ = self.wq_b(qr)
120 + q = q.view(-1, self.n_head, self.head_dim)
121 +
122 + q_pe, q_nope = torch.split(
123 + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
124 + )
125 + # Fused wk + weights_proj: one GEMM, then split.
126 + kw, _ = self.wk_weights_proj(hidden_states)
127 + k = kw[:, : self.head_dim]
128 + weights = kw[:, self.head_dim :]
129 +
130 + k = self.k_norm(k)
131 + k_pe, k_nope = torch.split(
132 + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
133 + )
134 +
135 + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
136 + # RoPE (NeoX) can introduce extra leading dims; reshape back to flat.
137 + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
138 + k_pe = k_pe.reshape(-1, 1, self.rope_dim)
139 +
140 + q = torch.cat([q_pe, q_nope], dim=-1)
141 + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
142 +
143 + # Only quant q here; k quant is fused with cache insertion.
144 + q = q.view(-1, self.head_dim)
145 + q_fp8, q_scale = per_token_group_quant_fp8(
146 + q,
147 + self.quant_block_size,
148 + column_major_scales=False,
149 + use_ue8m0=self.scale_fmt is not None,
150 + )
151 + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
152 + q_scale = q_scale.view(-1, self.n_head, 1)
153 +
154 + weights = (
155 + weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
156 + )
157 + weights = weights.squeeze(-1)
158 +
159 + return self.indexer_op(hidden_states, q_fp8, k, weights)
160 +
161 +
162 +class DeepseekV32Attention(MLAAttention):
163 + def __init__(
164 + self,
165 + vllm_config: VllmConfig,
166 + config: DeepseekV2Config | DeepseekV3Config,
167 + prefix: str,
168 + topk_indices_buffer: torch.Tensor | None = None,
169 + ) -> None:
170 + quant_config = vllm_config.quant_config
171 + cache_config = vllm_config.cache_config
172 +
173 + hidden_size = config.hidden_size
174 + qk_nope_head_dim = config.qk_nope_head_dim
175 + qk_rope_head_dim = config.qk_rope_head_dim
176 + v_head_dim = config.v_head_dim
177 + q_lora_rank = config.q_lora_rank
178 + kv_lora_rank = config.kv_lora_rank
179 + num_heads = config.num_attention_heads
180 +
181 + tp_size = get_tensor_model_parallel_world_size()
182 + assert num_heads % tp_size == 0
183 + num_local_heads = num_heads // tp_size
184 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
185 + scaling = qk_head_dim**-0.5
186 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
187 +
188 + # DSA checkpoints may use plain ("default") or yarn-scaled RoPE.
189 + if config.rope_parameters["rope_type"] != "default":
190 + config.rope_parameters["rope_type"] = (
191 + "deepseek_yarn"
192 + if config.rope_parameters.get("apply_yarn_scaling", True)
193 + else "deepseek_llama_scaling"
194 + )
195 + if config.rope_parameters["rope_type"] == "deepseek_yarn":
196 + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
197 + scaling_factor = config.rope_parameters["factor"]
198 + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
199 + scaling = scaling * mscale * mscale
200 +
201 + # DSA "shared indexer" pattern: only some layers carry an indexer; the
202 + # rest reuse the top-k written by the previous indexer layer into the
203 + # shared topk_indices_buffer. DeepSeek-V3.2 builds it on every layer
204 + # (index_topk_freq defaults to 1); GLM-5.2 uses index_topk_freq=4 so
205 + # only layers [0,1,2,6,10,...] (+ MTP) carry one.
206 + layer_id = extract_layer_index(prefix)
207 + index_topk_freq = getattr(config, "index_topk_freq", 1)
208 + index_topk_pattern = getattr(config, "index_topk_pattern", None)
209 + index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
210 + if index_topk_pattern is None:
211 + skip_topk = (
212 + max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0
213 + )
214 + elif 0 <= layer_id < len(index_topk_pattern):
215 + skip_topk = index_topk_pattern[layer_id] == "S"
216 + else:
217 + skip_topk = False
218 + # MTP/nextn layers always build a full indexer (they toggle at runtime).
219 + num_hidden_layers = getattr(config, "num_hidden_layers", None)
220 + is_mtp_layer = num_hidden_layers is not None and layer_id >= num_hidden_layers
221 +
222 + # Build kv_b_proj + indexer first; they are passed to MLAAttention.__init__
223 + # (which runs nn.Module.__init__ and registers them).
224 + kv_b_proj = ColumnParallelLinear(
225 + kv_lora_rank,
226 + num_heads * (qk_nope_head_dim + v_head_dim),
227 + bias=False,
228 + quant_config=quant_config,
229 + prefix=f"{prefix}.kv_b_proj",
230 + )
231 + indexer = None
232 + if not skip_topk or is_mtp_layer:
233 + indexer = DeepseekV32Indexer(
234 + vllm_config,
235 + config,
236 + hidden_size,
237 + q_lora_rank,
238 + quant_config,
239 + cache_config,
240 + topk_indices_buffer,
241 + prefix=f"{prefix}.indexer",
242 + )
243 +
244 + # Set up the MLA engine (impl, KV cache, scales, backend, registration,
245 + # and process_weights_after_loading) via the MLAAttention base.
246 + super().__init__(
247 + num_heads=num_local_heads,
248 + scale=scaling,
249 + qk_nope_head_dim=qk_nope_head_dim,
250 + qk_rope_head_dim=qk_rope_head_dim,
251 + v_head_dim=v_head_dim,
252 + q_lora_rank=q_lora_rank,
253 + kv_lora_rank=kv_lora_rank,
254 + kv_b_proj=kv_b_proj,
255 + cache_config=cache_config,
256 + quant_config=quant_config,
257 + prefix=f"{prefix}.attn",
258 + use_sparse=True,
259 + indexer=indexer,
260 + topk_indices_buffer=topk_indices_buffer,
261 + )
262 +
263 + self.num_local_heads = num_local_heads
264 + self.qk_head_dim = qk_head_dim
265 + self.indexer = indexer
266 + # Runtime toggle for index_share_for_mtp_iteration: MTP draft step 0
267 + # computes the top-k, steps 1+ set this True to reuse it.
268 + self.skip_topk = False
269 + # Whether the paged KV cache must be viewed as fp8 before the attention
270 + # (per-tensor fp8; the fp8_ds_mla layout is read as uint8).
271 + self._fp8_kv_needs_view = (
272 + is_quantized_kv_cache(self.kv_cache_dtype)
273 + and self.kv_cache_dtype != "fp8_ds_mla"
274 + )
275 + # Whether the backend takes an fp8-quantized query (FlashInfer sparse)
276 + # vs the (ql_nope, q_pe) tuple (FlashMLA sparse).
277 + self._use_concat_quant = (
278 + is_quantized_kv_cache(self.kv_cache_dtype)
279 + and self.impl.supports_quant_query_input
280 + )
281 +
282 + # Remaining MLA projections (registered on this module).
283 + self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
284 + hidden_size,
285 + [q_lora_rank, kv_lora_rank + qk_rope_head_dim],
286 + quant_config=quant_config,
287 + prefix=f"{prefix}.fused_qkv_a_proj",
288 + )
289 + self.q_a_layernorm = RMSNorm(q_lora_rank, eps=config.rms_norm_eps)
290 + self.q_b_proj = ColumnParallelLinear(
291 + q_lora_rank,
292 + num_heads * qk_head_dim,
293 + bias=False,
294 + quant_config=quant_config,
295 + prefix=f"{prefix}.q_b_proj",
296 + )
297 + self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps)
298 + self.o_proj = RowParallelLinear(
299 + num_heads * v_head_dim,
300 + hidden_size,
301 + bias=False,
302 + quant_config=quant_config,
303 + prefix=f"{prefix}.o_proj",
304 + )
305 +
306 + self.rotary_emb = get_rope(
307 + qk_rope_head_dim,
308 + max_position=max_position_embeddings,
309 + rope_parameters=config.rope_parameters,
310 + is_neox_style=False,
311 + )
312 + # Lightning indexer uses its own RoPE; interleave maps to non-NeoX.
313 + self.indexer_rope_emb = get_rope(
314 + qk_rope_head_dim,
315 + max_position=max_position_embeddings,
316 + rope_parameters=config.rope_parameters,
317 + is_neox_style=not getattr(config, "indexer_rope_interleave", False),
318 + )
319 +
320 + def forward( # type: ignore[override]
321 + self,
322 + positions: torch.Tensor,
323 + hidden_states: torch.Tensor,
324 + ) -> torch.Tensor:
325 + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
326 + q_c, kv_lora = qkv_lora.split(
327 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
328 + )
329 + q_c = self.q_a_layernorm(q_c)
330 + q = self.q_b_proj(q_c)[0]
331 +
332 + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
333 + kv_c_normed = self.kv_a_layernorm(kv_c)
334 +
335 + q = q.view(-1, self.num_local_heads, self.qk_head_dim)
336 + k_pe = k_pe.unsqueeze(1)
337 + q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
338 + positions, q[..., self.qk_nope_head_dim :], k_pe
339 + )
340 +
341 + num_tokens = hidden_states.shape[0]
342 + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
343 + q_nope = q_nope.transpose(0, 1) # (N, B, P)
344 + ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) # (B, N, L)
345 +
346 + # Lightning indexer writes the top-k indices into the shared buffer.
347 + # "Shared" layers (indexer is None) reuse the top-k from the previous
348 + # indexer layer already sitting in the buffer.
349 + if self.indexer is not None and not self.skip_topk:
350 + self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb) # type: ignore[operator]
351 +
352 + attn_latent = torch.empty(
353 + (num_tokens, self.num_local_heads, self.kv_lora_rank),
354 + dtype=q.dtype,
355 + device=q.device,
356 + )
357 + self._sparse_attention(kv_c_normed, k_pe, ql_nope, q_pe, attn_latent)
358 +
359 + # V up-projection + output projection are metadata-independent GEMMs and
360 + # stay captured.
361 + output = torch.empty(
362 + (num_tokens, self.num_local_heads * self.v_head_dim),
363 + dtype=q.dtype,
364 + device=q.device,
365 + )
366 + self._v_up_proj(attn_latent, out=output)
367 + return self.o_proj(output)[0]
368 +
369 + @eager_break_during_capture
370 + def _sparse_attention(
371 + self,
372 + kv_c_normed: torch.Tensor,
373 + k_pe: torch.Tensor,
374 + ql_nope: torch.Tensor,
375 + q_pe: torch.Tensor,
376 + attn_latent: torch.Tensor,
377 + ) -> None:
378 + forward_context = get_forward_context()
379 + attn_metadata_raw = forward_context.attn_metadata
380 + attn_metadata: MLACommonMetadata | None
381 + if isinstance(attn_metadata_raw, dict):
382 + attn_metadata = attn_metadata_raw[self.layer_name] # type: ignore[assignment]
383 + elif isinstance(attn_metadata_raw, list):
384 + # Speculative decoding: [0] is the base-model metadata dict.
385 + attn_metadata = attn_metadata_raw[0][self.layer_name] # type: ignore[assignment]
386 + else:
387 + attn_metadata = attn_metadata_raw
388 +
389 + slot_mapping = forward_context.slot_mapping
390 + assert isinstance(slot_mapping, dict)
391 + self.impl.do_kv_cache_update( # type: ignore[attr-defined]
392 + kv_c_normed,
393 + k_pe,
394 + self.kv_cache,
395 + slot_mapping.get(self.layer_name),
396 + self.kv_cache_dtype,
397 + self._k_scale,
398 + )
399 +
400 + if attn_metadata is None:
401 + # Profile / warmup: zero-fill for DP+EP determinism.
402 + attn_latent.zero_()
403 + return
404 +
405 + num_actual = attn_metadata.num_actual_tokens
406 + kv_cache = self.kv_cache
407 + if self._fp8_kv_needs_view:
408 + kv_cache = kv_cache.view(torch.float8_e4m3fn)
409 +
410 + ql_nope = ql_nope[:num_actual]
411 + q_pe = q_pe[:num_actual]
412 + # FlashInfer sparse takes a single fp8-quantized query; FlashMLA sparse
413 + # takes the (ql_nope, q_pe) tuple and concatenates internally.
414 + mqa_q: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
415 + if self._use_concat_quant:
416 + mqa_q = self._decode_concat_quant_fp8_op(ql_nope, q_pe, self._q_scale)
417 + else:
418 + mqa_q = (ql_nope, q_pe)
419 +
420 + attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined]
421 + attn_latent[:num_actual] = attn_out.view(
422 + num_actual, self.num_local_heads, self.kv_lora_rank
423 + )
vllm/models/deepseek_v32/nvidia/model.py new
+333
@@ -0,0 +1,333 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import typing
4 +from collections.abc import Callable, Iterable
5 +from itertools import islice
6 +
7 +import torch
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.distributed import get_pp_group
11 +from vllm.model_executor.layers.fused_moe import (
12 + fused_moe_make_expert_params_mapping,
13 +)
14 +from vllm.model_executor.layers.layernorm import RMSNorm
15 +from vllm.model_executor.layers.vocab_parallel_embedding import (
16 + VocabParallelEmbedding,
17 +)
18 +from vllm.model_executor.model_loader.weight_utils import (
19 + default_weight_loader,
20 + maybe_remap_kv_scale_name,
21 +)
22 +from vllm.model_executor.models.deepseek_v2 import (
23 + DeepseekV2ForCausalLM,
24 + DeepseekV2MLP,
25 + DeepseekV2MoE,
26 + _try_load_fp8_indexer_wk,
27 + get_spec_layer_idx_from_weight_name,
28 +)
29 +from vllm.model_executor.models.utils import (
30 + PPMissingLayer,
31 + get_pp_missing_layer_names,
32 + is_pp_missing_parameter,
33 + make_empty_intermediate_tensors_factory,
34 + make_layers,
35 +)
36 +from vllm.sequence import IntermediateTensors
37 +
38 +from .attention import DeepseekV32Attention
39 +
40 +
41 +class DeepseekV32DecoderLayer(torch.nn.Module):
42 + def __init__(
43 + self,
44 + vllm_config: VllmConfig,
45 + prefix: str,
46 + config=None,
47 + topk_indices_buffer: torch.Tensor | None = None,
48 + ) -> None:
49 + super().__init__()
50 +
51 + if config is None:
52 + config = vllm_config.model_config.hf_config
53 + quant_config = vllm_config.quant_config
54 + parallel_config = vllm_config.parallel_config
55 +
56 + self.hidden_size = config.hidden_size
57 + moe_layer_freq = getattr(config, "moe_layer_freq", 1)
58 + layer_idx = int(prefix.split(sep=".")[-1])
59 + self.layer_idx = layer_idx
60 + self.use_mha = False
61 +
62 + self.self_attn = DeepseekV32Attention(
63 + vllm_config=vllm_config,
64 + config=config,
65 + prefix=f"{prefix}.self_attn",
66 + topk_indices_buffer=topk_indices_buffer,
67 + )
68 +
69 + if (
70 + config.n_routed_experts is not None
71 + and layer_idx >= config.first_k_dense_replace
72 + and layer_idx % moe_layer_freq == 0
73 + ):
74 + self.mlp = DeepseekV2MoE(
75 + config=config,
76 + parallel_config=parallel_config,
77 + quant_config=quant_config,
78 + prefix=f"{prefix}.mlp",
79 + )
80 + else:
81 + self.mlp = DeepseekV2MLP(
82 + hidden_size=config.hidden_size,
83 + intermediate_size=config.intermediate_size,
84 + hidden_act=config.hidden_act,
85 + quant_config=quant_config,
86 + prefix=f"{prefix}.mlp",
87 + )
88 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
89 + self.post_attention_layernorm = RMSNorm(
90 + config.hidden_size, eps=config.rms_norm_eps
91 + )
92 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
93 +
94 + def forward(
95 + self,
96 + positions: torch.Tensor,
97 + hidden_states: torch.Tensor,
98 + residual: torch.Tensor | None,
99 + ) -> tuple[torch.Tensor, torch.Tensor]:
100 + if residual is None:
101 + residual = hidden_states
102 + hidden_states = self.input_layernorm(hidden_states)
103 + else:
104 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
105 + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
106 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
107 + hidden_states = self.mlp(hidden_states)
108 + return hidden_states, residual
109 +
110 +
111 +class DeepseekV32Model(torch.nn.Module):
112 + fall_back_to_pt_during_load = False
113 +
114 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
115 + super().__init__()
116 +
117 + config = vllm_config.model_config.hf_config
118 + quant_config = vllm_config.quant_config
119 + self.config = config
120 + from vllm.platforms import current_platform
121 +
122 + self.device = current_platform.device_type
123 +
124 + self.vocab_size = config.vocab_size
125 + # DSA is always sparse (has index_topk); allocate the shared top-k
126 + # buffer the indexer writes and the sparse MLA backend reads.
127 + self.is_v32 = True
128 + topk_indices_buffer = torch.empty(
129 + vllm_config.scheduler_config.max_num_batched_tokens,
130 + config.index_topk,
131 + dtype=torch.int32,
132 + device=self.device,
133 + )
134 +
135 + if get_pp_group().is_first_rank:
136 + self.embed_tokens = VocabParallelEmbedding(
137 + config.vocab_size,
138 + config.hidden_size,
139 + quant_config=quant_config,
140 + prefix=f"{prefix}.embed_tokens",
141 + )
142 + else:
143 + self.embed_tokens = PPMissingLayer()
144 +
145 + self.start_layer, self.end_layer, self.layers = make_layers(
146 + config.num_hidden_layers,
147 + lambda prefix: DeepseekV32DecoderLayer(
148 + vllm_config=vllm_config,
149 + prefix=prefix,
150 + topk_indices_buffer=topk_indices_buffer,
151 + ),
152 + prefix=f"{prefix}.layers",
153 + )
154 +
155 + if get_pp_group().is_last_rank:
156 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
157 + else:
158 + self.norm = PPMissingLayer()
159 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
160 + ["hidden_states", "residual"], config.hidden_size
161 + )
162 +
163 + self.aux_hidden_state_layers = tuple[int, ...]()
164 + self.num_redundant_experts = (
165 + vllm_config.parallel_config.eplb_config.num_redundant_experts
166 + )
167 +
168 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
169 + return self.embed_tokens(input_ids)
170 +
171 + def forward(
172 + self,
173 + input_ids: torch.Tensor | None,
174 + positions: torch.Tensor,
175 + intermediate_tensors: IntermediateTensors | None = None,
176 + inputs_embeds: torch.Tensor | None = None,
177 + ) -> torch.Tensor | IntermediateTensors:
178 + if get_pp_group().is_first_rank:
179 + if inputs_embeds is not None:
180 + hidden_states = inputs_embeds
181 + else:
182 + assert input_ids is not None
183 + hidden_states = self.embed_input_ids(input_ids)
184 + residual = None
185 + else:
186 + assert intermediate_tensors is not None
187 + hidden_states = intermediate_tensors["hidden_states"]
188 + residual = intermediate_tensors["residual"]
189 +
190 + aux_hidden_states = []
191 + for idx, layer in enumerate(
192 + islice(self.layers, self.start_layer, self.end_layer),
193 + start=self.start_layer,
194 + ):
195 + if idx in self.aux_hidden_state_layers:
196 + aux_hidden_states.append(hidden_states + residual)
197 + hidden_states, residual = layer(positions, hidden_states, residual)
198 +
199 + if not get_pp_group().is_last_rank:
200 + return IntermediateTensors(
201 + {"hidden_states": hidden_states, "residual": residual}
202 + )
203 +
204 + hidden_states, _ = self.norm(hidden_states, residual)
205 + if len(aux_hidden_states) > 0:
206 + return hidden_states, aux_hidden_states
207 + return hidden_states
208 +
209 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
210 + # DSA-only: MLA (fused_qkv_a_proj) + the fused indexer wk/weights_proj +
211 + # routed experts. No MHA (qkv_proj) or ROCm shared-expert-fusion paths.
212 + stacked_params_mapping = [
213 + ("gate_up_proj", "gate_proj", 0),
214 + ("gate_up_proj", "up_proj", 1),
215 + ("fused_qkv_a_proj", "q_a_proj", 0),
216 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
217 + ("wk_weights_proj", "wk", 0),
218 + ("wk_weights_proj", "weights_proj", 1),
219 + ]
220 + expert_params_mapping = fused_moe_make_expert_params_mapping(
221 + self,
222 + ckpt_gate_proj_name="gate_proj",
223 + ckpt_down_proj_name="down_proj",
224 + ckpt_up_proj_name="up_proj",
225 + num_experts=self.config.n_routed_experts,
226 + num_redundant_experts=self.num_redundant_experts,
227 + )
228 +
229 + pp_missing_layer_names = get_pp_missing_layer_names(self)
230 + params_dict = dict(self.named_parameters())
231 + loaded_params: set[str] = set()
232 + _pending_wk_fp8: dict = {}
233 + for name, loaded_weight in weights:
234 + if "rotary_emb.inv_freq" in name:
235 + continue
236 + # MTP / nextn layers are loaded by the MTP model, not here.
237 + if get_spec_layer_idx_from_weight_name(self.config, name) is not None:
238 + continue
239 + if _try_load_fp8_indexer_wk(
240 + name,
241 + loaded_weight,
242 + _pending_wk_fp8,
243 + params_dict,
244 + loaded_params,
245 + pp_missing_layer_names,
246 + ):
247 + continue
248 +
249 + for param_name, weight_name, shard_id in stacked_params_mapping:
250 + if weight_name not in name:
251 + continue
252 + # Experts are handled below; skip here before the name rewrite.
253 + if ("mlp.experts." in name) and name not in params_dict:
254 + continue
255 + name_mapped = name.replace(weight_name, param_name)
256 + if (
257 + param_name == "fused_qkv_a_proj"
258 + ) and name_mapped not in params_dict:
259 + continue
260 + name = name_mapped
261 + if name.endswith(".bias") and name not in params_dict:
262 + continue
263 + if is_pp_missing_parameter(name, self):
264 + continue
265 + param = params_dict[name]
266 + param.weight_loader(param, loaded_weight, shard_id)
267 + break
268 + else:
269 + is_expert_weight = False
270 + for mapping in expert_params_mapping:
271 + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
272 + if weight_name not in name:
273 + continue
274 + is_expert_weight = True
275 + name_mapped = name.replace(weight_name, param_name)
276 + if is_pp_missing_parameter(name_mapped, self):
277 + continue
278 + param = params_dict[name_mapped]
279 + weight_loader = typing.cast(
280 + Callable[..., bool], param.weight_loader
281 + )
282 + success = weight_loader(
283 + param,
284 + loaded_weight,
285 + name_mapped,
286 + shard_id=shard_id,
287 + expert_id=expert_id,
288 + return_success=True,
289 + )
290 + if success:
291 + name = name_mapped
292 + break
293 + else:
294 + if is_expert_weight:
295 + continue
296 + if name.endswith(".bias") and name not in params_dict:
297 + continue
298 + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
299 + if name is None:
300 + continue
301 + if is_pp_missing_parameter(name, self):
302 + continue
303 + param = params_dict[name]
304 + loader = getattr(param, "weight_loader", default_weight_loader)
305 + loader(param, loaded_weight)
306 + loaded_params.add(name)
307 + return loaded_params
308 +
309 +
310 +class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
311 + """DSA causal LM — DeepSeek V2/V3 orchestration with the DSA backbone.
312 +
313 + Serves DeepSeek V3.2 and any architecture reusing DSA (e.g. GLM-5.2).
314 + """
315 +
316 + model_cls = DeepseekV32Model
317 +
318 + def set_moe_parameters(self):
319 + # Same as the base, but keyed on the MoE block type rather than the
320 + # decoder-layer type (DeepseekV32DecoderLayer is a plain nn.Module).
321 + self.expert_weights = []
322 + self.num_expert_groups = getattr(self.config, "n_group", 1)
323 + self.moe_layers = []
324 + self.moe_mlp_layers = []
325 + example_moe = None
326 + for layer in self.model.layers:
327 + if isinstance(layer, PPMissingLayer):
328 + continue
329 + if isinstance(layer.mlp, DeepseekV2MoE):
330 + example_moe = layer.mlp
331 + self.moe_mlp_layers.append(layer.mlp)
332 + self.moe_layers.append(layer.mlp.experts)
333 + self.extract_moe_parameters(example_moe)
vllm/models/deepseek_v32/nvidia/mtp.py new
+390
@@ -0,0 +1,390 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import typing
5 +from collections.abc import Callable, Iterable
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm._aiter_ops import rocm_aiter_ops
11 +from vllm.config import VllmConfig
12 +from vllm.model_executor.layers.fused_moe import (
13 + fused_moe_make_expert_params_mapping,
14 +)
15 +from vllm.model_executor.layers.layernorm import RMSNorm
16 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
17 +from vllm.model_executor.layers.vocab_parallel_embedding import (
18 + VocabParallelEmbedding,
19 +)
20 +from vllm.model_executor.model_loader.weight_utils import (
21 + default_weight_loader,
22 + maybe_remap_kv_scale_name,
23 +)
24 +from vllm.model_executor.models.deepseek_mtp import SharedHead
25 +from vllm.model_executor.models.deepseek_v2 import (
26 + DeepseekV2MixtureOfExperts,
27 + DeepseekV2MoE,
28 + _try_load_fp8_indexer_wk,
29 + get_spec_layer_idx_from_weight_name,
30 +)
31 +from vllm.model_executor.models.utils import (
32 + get_pp_missing_layer_names,
33 + maybe_prefix,
34 +)
35 +from vllm.platforms import current_platform
36 +from vllm.sequence import IntermediateTensors
37 +
38 +from .model import DeepseekV32DecoderLayer
39 +
40 +
41 +class DeepseekV32MultiTokenPredictorLayer(nn.Module):
42 + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
43 + super().__init__()
44 + assert vllm_config.speculative_config is not None
45 + config = vllm_config.speculative_config.draft_model_config.hf_config
46 + self.config = config
47 + quant_config = vllm_config.quant_config
48 +
49 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
50 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
51 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
52 +
53 + topk_indices_buffer = torch.empty(
54 + vllm_config.scheduler_config.max_num_batched_tokens,
55 + config.index_topk,
56 + dtype=torch.int32,
57 + device=current_platform.device_type,
58 + )
59 + self.shared_head = SharedHead(
60 + config=config, prefix=prefix, quant_config=quant_config
61 + )
62 + self.mtp_block = DeepseekV32DecoderLayer(
63 + vllm_config,
64 + prefix,
65 + config=config,
66 + topk_indices_buffer=topk_indices_buffer,
67 + )
68 +
69 + def forward(
70 + self,
71 + input_ids: torch.Tensor,
72 + positions: torch.Tensor,
73 + previous_hidden_states: torch.Tensor,
74 + inputs_embeds: torch.Tensor | None = None,
75 + spec_step_index: int = 0,
76 + ) -> torch.Tensor:
77 + assert inputs_embeds is not None
78 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
79 + inputs_embeds = self.enorm(inputs_embeds)
80 + previous_hidden_states = self.hnorm(previous_hidden_states)
81 + hidden_states = self.eh_proj(
82 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
83 + )
84 + hidden_states, residual = self.mtp_block(
85 + positions=positions, hidden_states=hidden_states, residual=None
86 + )
87 + # Return the pre-final-norm recycle hidden (re-fed as the next spec
88 + # step's previous_hidden_states); shared_head norm is applied in
89 + # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract.
90 + return residual + hidden_states
91 +
92 +
93 +class DeepseekV32MultiTokenPredictor(nn.Module):
94 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
95 + super().__init__()
96 + config = vllm_config.model_config.hf_config
97 + self.mtp_start_layer_idx = config.num_hidden_layers
98 + self.num_mtp_layers = config.num_nextn_predict_layers
99 + self.layers = torch.nn.ModuleDict(
100 + {
101 + str(idx): DeepseekV32MultiTokenPredictorLayer(
102 + vllm_config, f"{prefix}.layers.{idx}"
103 + )
104 + for idx in range(
105 + self.mtp_start_layer_idx,
106 + self.mtp_start_layer_idx + self.num_mtp_layers,
107 + )
108 + }
109 + )
110 + self.embed_tokens = VocabParallelEmbedding(
111 + config.vocab_size,
112 + config.hidden_size,
113 + prefix=maybe_prefix(prefix, "embed_tokens"),
114 + )
115 + self.logits_processor = LogitsProcessor(config.vocab_size)
116 +
117 + def set_skip_topk(self, skip: bool):
118 + # index_share_for_mtp_iteration: step 0 computes top-k, steps 1+ reuse.
119 + for layer in self.layers.values():
120 + self_attn = getattr(layer.mtp_block, "self_attn", None)
121 + if self_attn is not None and hasattr(self_attn, "skip_topk"):
122 + self_attn.skip_topk = skip
123 +
124 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
125 + return self.embed_tokens(input_ids)
126 +
127 + def forward(
128 + self,
129 + input_ids: torch.Tensor,
130 + positions: torch.Tensor,
131 + previous_hidden_states: torch.Tensor,
132 + inputs_embeds: torch.Tensor | None = None,
133 + spec_step_idx: int = 0,
134 + ) -> torch.Tensor:
135 + if inputs_embeds is None:
136 + inputs_embeds = self.embed_tokens(input_ids)
137 + current_step_idx = spec_step_idx % self.num_mtp_layers
138 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
139 + input_ids,
140 + positions,
141 + previous_hidden_states,
142 + inputs_embeds,
143 + current_step_idx,
144 + )
145 +
146 + def compute_logits(
147 + self,
148 + hidden_states: torch.Tensor,
149 + spec_step_idx: int = 0,
150 + ) -> torch.Tensor:
151 + current_step_idx = spec_step_idx % self.num_mtp_layers
152 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
153 + return self.logits_processor(
154 + mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
155 + )
156 +
157 +
158 +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
159 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
160 + super().__init__()
161 + self.config = vllm_config.model_config.hf_config
162 + self.quant_config = vllm_config.quant_config
163 + self.model = DeepseekV32MultiTokenPredictor(
164 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
165 + )
166 + self.set_moe_parameters()
167 +
168 + def set_moe_parameters(self):
169 + self.expert_weights = []
170 + self.num_moe_layers = self.config.num_nextn_predict_layers
171 + self.num_expert_groups = self.config.n_group
172 + self.moe_layers = []
173 + self.moe_mlp_layers = []
174 + example_moe = None
175 + for layer in self.model.layers.values():
176 + mlp = layer.mtp_block.mlp
177 + if isinstance(mlp, DeepseekV2MoE):
178 + example_moe = mlp
179 + self.moe_mlp_layers.append(mlp)
180 + self.moe_layers.append(mlp.experts)
181 + self.extract_moe_parameters(example_moe)
182 +
183 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
184 + return self.model.embed_input_ids(input_ids)
185 +
186 + def forward(
187 + self,
188 + input_ids: torch.Tensor | None,
189 + positions: torch.Tensor,
190 + hidden_states: torch.Tensor,
191 + intermediate_tensors: IntermediateTensors | None = None,
192 + inputs_embeds: torch.Tensor | None = None,
193 + spec_step_idx: int = 0,
194 + ) -> torch.Tensor:
195 + return self.model(
196 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
197 + )
198 +
199 + def compute_logits(
200 + self,
201 + hidden_states: torch.Tensor,
202 + spec_step_idx: int = 0,
203 + ) -> torch.Tensor | None:
204 + return self.model.compute_logits(hidden_states, spec_step_idx)
205 +
206 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
207 + spec_layer_weight_names = [
208 + "embed_tokens",
209 + "enorm",
210 + "hnorm",
211 + "eh_proj",
212 + "shared_head",
213 + ]
214 + shared_weight_names = ["embed_tokens"]
215 + spec_layer_weight = False
216 + shared_weight = False
217 + for weight_name in spec_layer_weight_names:
218 + if weight_name in name:
219 + spec_layer_weight = True
220 + if weight_name in shared_weight_names:
221 + shared_weight = True
222 + break
223 + if not spec_layer_weight:
224 + name = name.replace(
225 + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
226 + )
227 + elif shared_weight:
228 + name = name.replace(f"model.layers.{spec_layer}.", "model.")
229 + return name
230 +
231 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
232 + rocm_aiter_moe_shared_expert_enabled = (
233 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
234 + )
235 + stacked_params_mapping = [
236 + ("gate_up_proj", "gate_proj", 0),
237 + ("gate_up_proj", "up_proj", 1),
238 + ("fused_qkv_a_proj", "q_a_proj", 0),
239 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
240 + ("wk_weights_proj", "wk", 0),
241 + ("wk_weights_proj", "weights_proj", 1),
242 + ]
243 + expert_params_mapping = fused_moe_make_expert_params_mapping(
244 + self,
245 + ckpt_gate_proj_name="gate_proj",
246 + ckpt_down_proj_name="down_proj",
247 + ckpt_up_proj_name="up_proj",
248 + num_experts=self.config.n_routed_experts
249 + + (
250 + self.config.n_shared_experts
251 + if rocm_aiter_moe_shared_expert_enabled
252 + else 0
253 + ),
254 + )
255 +
256 + pp_missing_layer_names = get_pp_missing_layer_names(self)
257 + params_dict = dict(self.named_parameters())
258 + loaded_params: set[str] = set()
259 + _pending_wk_fp8: dict = {}
260 + for name, loaded_weight in weights:
261 + if "rotary_emb.inv_freq" in name:
262 + continue
263 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
264 + if spec_layer is None:
265 + continue
266 + is_fusion_moe_shared_experts_layer = (
267 + rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
268 + )
269 + name = self._rewrite_spec_layer_name(spec_layer, name)
270 +
271 + if _try_load_fp8_indexer_wk(
272 + name,
273 + loaded_weight,
274 + _pending_wk_fp8,
275 + params_dict,
276 + loaded_params,
277 + pp_missing_layer_names,
278 + ):
279 + continue
280 +
281 + for param_name, weight_name, shard_id in stacked_params_mapping:
282 + if weight_name not in name:
283 + continue
284 + if ("mlp.experts." in name) and name not in params_dict:
285 + continue
286 + if is_fusion_moe_shared_experts_layer:
287 + continue
288 + name_mapped = name.replace(weight_name, param_name)
289 + if (
290 + param_name == "fused_qkv_a_proj"
291 + ) and name_mapped not in params_dict:
292 + continue
293 + else:
294 + name = name_mapped
295 + if name.endswith(".bias") and name not in params_dict:
296 + continue
297 + param = params_dict[name]
298 + weight_loader = param.weight_loader
299 + weight_loader(param, loaded_weight, shard_id)
300 + break
301 + else:
302 + num_chunks = 1
303 + if is_fusion_moe_shared_experts_layer:
304 + num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
305 + split_dim = (
306 + 1
307 + if ("down_proj.weight" in name and loaded_weight.ndim > 1)
308 + else 0
309 + )
310 + total = loaded_weight.shape[split_dim]
311 + assert total % num_chunks == 0
312 + chunk_size = total // num_chunks
313 +
314 + for j in range(num_chunks):
315 + chunk_name = name
316 + weight_to_load = loaded_weight
317 + if is_fusion_moe_shared_experts_layer:
318 + chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
319 + if loaded_weight.ndim == 1:
320 + weight_to_load = loaded_weight[chunk_slice]
321 + elif split_dim == 0:
322 + weight_to_load = loaded_weight[chunk_slice, :]
323 + else:
324 + weight_to_load = loaded_weight[:, chunk_slice]
325 + chunk_name = name.replace(
326 + "mlp.shared_experts",
327 + f"mlp.experts.{self.config.n_routed_experts + j}",
328 + )
329 +
330 + is_expert_weight = False
331 + for mapping in expert_params_mapping:
332 + param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
333 + if weight_name not in chunk_name:
334 + continue
335 + is_expert_weight = True
336 + name_mapped = chunk_name.replace(weight_name, param_name)
337 + param = params_dict[name_mapped]
338 + weight_loader = typing.cast(
339 + Callable[..., bool], param.weight_loader
340 + )
341 + success = weight_loader(
342 + param,
343 + weight_to_load,
344 + name_mapped,
345 + shard_id=shard_id,
346 + expert_id=expert_id,
347 + return_success=True,
348 + )
349 + if success:
350 + if not is_fusion_moe_shared_experts_layer:
351 + name = name_mapped
352 + else:
353 + loaded_params.add(name_mapped)
354 + break
355 + else:
356 + if is_expert_weight:
357 + continue
358 + if name.endswith(".bias") and name not in params_dict:
359 + continue
360 + name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
361 + if name is None:
362 + continue
363 + if (
364 + spec_layer != self.model.mtp_start_layer_idx
365 + and ".layers" not in name
366 + ):
367 + continue
368 + param = params_dict[name]
369 + weight_loader = getattr(
370 + param, "weight_loader", default_weight_loader
371 + )
372 + weight_loader(param, loaded_weight)
373 + if not is_fusion_moe_shared_experts_layer:
374 + loaded_params.add(name)
375 +
376 + loaded_layers: set[int] = set()
377 + for param_name in loaded_params:
378 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
379 + if spec_layer is not None:
380 + loaded_layers.add(spec_layer)
381 + for layer_idx in range(
382 + self.model.mtp_start_layer_idx,
383 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
384 + ):
385 + if layer_idx not in loaded_layers:
386 + raise ValueError(
387 + f"MTP speculative decoding layer {layer_idx} weights "
388 + f"missing from checkpoint."
389 + )
390 + return loaded_params