A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 05:30 UTC
[Spec Decode] Support non-MTP speculation for NemotronH (#43130)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
May 20, 2026 at 09:15 UTC
0a508743d42a26786c1432bb7f2e93f8111b6383
1 file changed
+16
-2
vllm/model_executor/models/nemotron_h.py
+16
-2
@@ -64,9 +64,12 @@ from vllm.model_executor.model_loader.weight_utils import (
64
maybe_remap_kv_scale_name,
65
)
66
from vllm.model_executor.models.interfaces import (
67
+ EagleModelMixin,
68
HasInnerState,
69
IsHybrid,
70
MixtureOfExperts,
71
+ SupportsEagle,
72
+ SupportsEagle3,
73
SupportsLoRA,
74
SupportsMambaPrefixCaching,
75
SupportsPP,
@@ -539,7 +542,7 @@ ALL_DECODER_LAYER_TYPES = {
542
543
544
@support_torch_compile
542
-class NemotronHModel(nn.Module):
545
+class NemotronHModel(nn.Module, EagleModelMixin):
546
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
547
super().__init__()
548
@@ -605,18 +608,27 @@ class NemotronHModel(nn.Module):
608
hidden_states = intermediate_tensors["hidden_states"]
609
residual = intermediate_tensors["residual"]
610
608
- for layer in islice(self.layers, self.start_layer, self.end_layer):
611
+ aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
612
+ for idx, layer in enumerate(
613
+ islice(self.layers, self.start_layer, self.end_layer)
614
+ ):
615
hidden_states, residual = layer(
616
positions=positions,
617
hidden_states=hidden_states,
618
residual=residual,
619
)
620
+ self._maybe_add_hidden_state(
621
+ aux_hidden_states, idx + 1, hidden_states, residual
622
+ )
623
624
if not get_pp_group().is_last_rank:
625
return IntermediateTensors(
626
{"hidden_states": hidden_states, "residual": residual}
627
)
628
hidden_states, _ = self.norm_f(hidden_states, residual)
629
+
630
+ if len(aux_hidden_states) > 0:
631
+ return hidden_states, aux_hidden_states
632
return hidden_states
633
634
def is_spec_layer(self, config: NemotronHConfig, weight_name: str) -> bool:
@@ -767,6 +779,8 @@ class NemotronHForCausalLM(
779
HasInnerState,
780
SupportsLoRA,
781
SupportsPP,
782
+ SupportsEagle,
783
+ SupportsEagle3,
784
IsHybrid,
785
SupportsQuant,
786
MixtureOfExperts,