[Feature][Spec Decode] MTP with separate (possibly quantized) lm head for nemotron (#54574)
Signed-off-by: Yoav Miron <yomiron@nvidia.com> Signed-off-by: YoavMiron <yomiron@nvidia.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
YoavMiron committed
Sep 10, 2026 at 15:26 UTC
40e6042ec83eb8f2971f21043a5da40496bd188a
2 files changed
+60
-22
vllm/config/speculative.py
+2
-3
@@ -620,16 +620,15 @@ class SpeculativeConfig:
620
)
621
factors.append(uses_aux_hidden_states)
622
623
- if uses_aux_hidden_states and self.draft_model_config is not None:
623
+ if self.draft_model_config is not None:
624
factors.append(self.draft_model_config.compute_hash())
625
-
625
# The specific layers used also affect the computation graph.
626
layer_ids = getattr(
627
self.draft_model_config.hf_config,
628
"eagle_aux_hidden_state_layer_ids",
629
None,
630
)
632
- if layer_ids is not None:
631
+ if layer_ids is not None and uses_aux_hidden_states:
632
# Convert to tuple to make it hashable
633
factors.append(tuple(layer_ids))
634
vllm/model_executor/models/nemotron_h_mtp.py
+58
-19
@@ -22,16 +22,20 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
22
ParallelLMHead,
23
VocabParallelEmbedding,
24
)
25
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
+from vllm.model_executor.model_loader.weight_utils import (
26
+ default_weight_loader,
27
+ maybe_remap_kv_scale_name,
28
+)
29
from vllm.model_executor.models.utils import (
30
WeightsMapper,
31
+ get_draft_quant_config,
32
make_empty_intermediate_tensors_factory,
33
maybe_prefix,
34
)
35
from vllm.sequence import IntermediateTensors
36
from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
37
34
-from .interfaces import SupportsPP
38
+from .interfaces import SupportsPP, SupportsQuant
39
from .nemotron_h import (
40
NemotronHAttentionDecoderLayer,
41
NemotronHMoEDecoderLayer,
@@ -214,10 +218,22 @@ class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer):
218
class NemotronHMultiTokenPredictor(nn.Module):
219
"""MTP predictor with NemotronH layers."""
220
217
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
221
+ def __init__(
222
+ self,
223
+ *,
224
+ vllm_config: VllmConfig,
225
+ quant_config: QuantizationConfig | None = None,
226
+ prefix: str = "",
227
+ ):
228
super().__init__()
229
220
- config = vllm_config.model_config.hf_config.get_text_config()
230
+ speculative_config = vllm_config.speculative_config
231
+ assert speculative_config is not None
232
+ draft_model_config = speculative_config.draft_model_config
233
+ assert draft_model_config is not None
234
+ config = draft_model_config.hf_config.get_text_config()
235
+ if quant_config is None:
236
+ quant_config = get_draft_quant_config(vllm_config)
237
238
self.config = config
239
self.vocab_size = config.vocab_size
@@ -257,9 +273,9 @@ class NemotronHMultiTokenPredictor(nn.Module):
273
common_kwargs = dict(
274
config=config,
275
layer_idx=self.mtp_start_layer_idx + i,
260
- model_config=vllm_config.model_config,
276
+ model_config=draft_model_config,
277
cache_config=vllm_config.cache_config,
262
- quant_config=vllm_config.quant_config,
278
+ quant_config=quant_config,
279
parallel_config=vllm_config.parallel_config,
280
prefix=layer_prefix,
281
has_start_projections=is_start_of_step,
@@ -310,14 +326,9 @@ class NemotronHMultiTokenPredictor(nn.Module):
326
return hidden_states
327
328
313
-class NemotronHMTP(nn.Module, SupportsPP):
329
+class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):
330
"""NemotronH MTP model."""
331
316
- # Quant configs name modules in checkpoint space ("language_model.mtp.layers.0*"),
317
- # but this draft is built under "mtp" (maybe_prefix below). SupportsQuant only
318
- # re-roots exclude_modules when the model defines a mapper, so without one the
319
- # exclusions never match and the MTP experts are wrongly quantized. Mirrors the
320
- # prefix handling load_weights() already does by hand.
332
hf_to_vllm_mapper = WeightsMapper(
333
orig_to_new_prefix={"language_model.": ""},
334
orig_to_new_substr={"embeddings": "embed_tokens"},
@@ -333,11 +344,13 @@ class NemotronHMTP(nn.Module, SupportsPP):
344
345
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
346
super().__init__()
336
- config = vllm_config.model_config.hf_config.get_text_config()
347
+ speculative_config = vllm_config.speculative_config
348
+ assert speculative_config is not None
349
+ draft_model_config = speculative_config.draft_model_config
350
+ assert draft_model_config is not None
351
+ config = draft_model_config.hf_config.get_text_config()
352
self.vllm_config = vllm_config
353
self.config = config
339
- self.quant_config = vllm_config.quant_config
340
-
354
# Needed for load_weights mapping
355
self.mtp_start_layer_idx = config.num_hidden_layers
356
@@ -350,13 +363,17 @@ class NemotronHMTP(nn.Module, SupportsPP):
363
364
# MTP predictor
365
self.model = NemotronHMultiTokenPredictor(
353
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
366
+ vllm_config=vllm_config,
367
+ quant_config=self.quant_config,
368
+ prefix=maybe_prefix(prefix, "mtp"),
369
)
370
371
# LM head for generating logits
372
+ self.has_own_lm_head = False
373
self.lm_head = ParallelLMHead(
374
self.config.vocab_size,
375
self.config.hidden_size,
376
+ quant_config=self.quant_config,
377
prefix=maybe_prefix(prefix, "lm_head"),
378
)
379
@@ -366,6 +383,11 @@ class NemotronHMTP(nn.Module, SupportsPP):
383
self.model.make_empty_intermediate_tensors
384
)
385
386
+ @staticmethod
387
+ def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
388
+ vllm_config = kwargs.get("vllm_config")
389
+ return get_draft_quant_config(vllm_config)
390
+
391
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
392
return self.model.get_input_embeddings(input_ids)
393
@@ -428,9 +450,16 @@ class NemotronHMTP(nn.Module, SupportsPP):
450
# MTP weights are nested in "language_model."
451
# in Multimodal Nemotron-H checkpoints.
452
name = name.removeprefix("language_model.")
431
-
432
- # Only process MTP weights - skip all non-MTP weights
433
- if not name.startswith("mtp.") and "embeddings" not in name:
453
+ is_lm_head_weight = name.startswith("lm_head.")
454
+ if is_lm_head_weight:
455
+ self.has_own_lm_head = True
456
+ # Only process MTP and LM head weights -
457
+ # skip all non-MTP and non-LM head weights
458
+ if (
459
+ not name.startswith("mtp.")
460
+ and "embeddings" not in name
461
+ and not is_lm_head_weight
462
+ ):
463
continue
464
# Skip rotary embeddings (computed, not loaded)
465
if "rotary_emb.inv_freq" in name:
@@ -443,6 +472,11 @@ class NemotronHMTP(nn.Module, SupportsPP):
472
if name.startswith("backbone."):
473
name = name.replace("backbone.", "model.")
474
475
+ if "scale" in name or "zero_point" in name:
476
+ name = maybe_remap_kv_scale_name(name, params_dict)
477
+ if name is None:
478
+ continue
479
+
480
# Handle stacked parameters (qkv_proj) for attention layers
481
is_stacked = False
482
for param_name, weight_name, shard_id in stacked_params_mapping:
@@ -516,4 +550,9 @@ class NemotronHMTP(nn.Module, SupportsPP):
550
weight_loader(param, loaded_weight)
551
loaded_params.add(name)
552
553
+ if not self.has_own_lm_head:
554
+ loaded_params.update(
555
+ name for name in params_dict if name.startswith("lm_head.")
556
+ )
557
+
558
return loaded_params