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 9, 13:33 UTC
Support heterogeneous NemotronHPuzzle model (#32549)
Signed-off-by: <dafrimi@nvidia.com> Signed-off-by: Daniel Afrimi <dafrimi@nvidia.com> Signed-off-by: root <dafrimi@nvidia.com>
danielafrimi committed
Jan 27, 2026 at 17:55 UTC
83fb2d09e8f6f70e99c68521b5a0218119cf5cfb
5 files changed
+75
-5
tests/models/registry.py
+5
@@ -405,6 +405,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
405
"NemotronHForCausalLM": _HfExamplesInfo(
406
"nvidia/Nemotron-H-8B-Base-8K", trust_remote_code=True
407
),
408
+ "NemotronHPuzzleForCausalLM": _HfExamplesInfo(
409
+ "",
410
+ trust_remote_code=True,
411
+ is_available_online=False,
412
+ ),
413
"OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
414
"Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
415
"Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"),
vllm/model_executor/models/config.py
+1
@@ -603,4 +603,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
603
"FalconMambaForCausalLM": MambaModelConfig,
604
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
605
"NemotronHForCausalLM": NemotronHForCausalLMConfig,
606
+ "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
607
}
vllm/model_executor/models/nemotron_h.py
+43
-3
@@ -354,8 +354,12 @@ class NemotronHMoEDecoderLayer(nn.Module):
354
super().__init__()
355
self.config = config
356
357
+ # Get per-layer config for heterogeneous models if exsist
358
+ get_layer_config = getattr(config, "get_nemotron_h_config_for_layer", None)
359
+ layer_config = get_layer_config(layer_idx) if get_layer_config else config
360
+
361
self.mixer = NemotronHMoE(
358
- config,
362
+ layer_config,
363
quant_config=quant_config,
364
parallel_config=parallel_config,
365
prefix=f"{prefix}.mixer",
@@ -479,6 +483,9 @@ class NemotronHAttention(nn.Module):
483
prefix=f"{prefix}.o_proj",
484
)
485
486
+ # Get per-layer sliding window from config (for heterogeneous models)
487
+ sliding_window = getattr(config, "sliding_window", None)
488
+
489
self.attn = Attention(
490
self.num_heads,
491
self.head_dim,
@@ -487,6 +494,7 @@ class NemotronHAttention(nn.Module):
494
cache_config=cache_config,
495
quant_config=quant_config,
496
prefix=f"{prefix}.attn",
497
+ per_layer_sliding_window=sliding_window,
498
)
499
500
def forward(
@@ -514,8 +522,12 @@ class NemotronHAttentionDecoderLayer(nn.Module):
522
) -> None:
523
super().__init__()
524
525
+ # Get per-layer config for heterogeneous models if exsist
526
+ get_layer_config = getattr(config, "get_nemotron_h_config_for_layer", None)
527
+ layer_config = get_layer_config(layer_idx) if get_layer_config else config
528
+
529
self.mixer = NemotronHAttention(
518
- config,
530
+ layer_config,
531
layer_idx,
532
model_config,
533
cache_config,
@@ -631,6 +643,34 @@ class NemotronHModel(nn.Module):
643
hidden_states, _ = self.norm_f(hidden_states, residual)
644
return hidden_states
645
646
+ def _get_max_n_routed_experts(self) -> int:
647
+ """Get max n_routed_experts from config or block_configs for puzzle models.
648
+
649
+ For heterogeneous models with varying expert counts per layer,
650
+ returns the MAX to ensure all expert weights can be loaded.
651
+ """
652
+ # First try top-level attribute
653
+ n_routed_experts = getattr(self.config, "n_routed_experts", None)
654
+ if n_routed_experts is not None:
655
+ return n_routed_experts
656
+
657
+ # For puzzle models, get MAX from all MoE blocks in block_configs
658
+ # (different layers may have different expert counts)
659
+ max_experts = 0
660
+ block_configs = getattr(self.config, "block_configs", None)
661
+ if block_configs:
662
+ for block in block_configs:
663
+ if isinstance(block, dict):
664
+ if block.get("block_type") == "moe":
665
+ max_experts = max(max_experts, block.get("n_routed_experts", 0))
666
+ else:
667
+ # HF converts dicts to objects with attributes
668
+ if getattr(block, "block_type", "") == "moe":
669
+ max_experts = max(
670
+ max_experts, getattr(block, "n_routed_experts", 0)
671
+ )
672
+ return max_experts
673
+
674
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
675
if self.has_moe:
676
# (param_name, weight_name, expert_id, shard_id)
@@ -643,7 +683,7 @@ class NemotronHModel(nn.Module):
683
ckpt_gate_proj_name="up_proj",
684
ckpt_down_proj_name="down_proj",
685
ckpt_up_proj_name="",
646
- num_experts=self.config.n_routed_experts,
686
+ num_experts=self._get_max_n_routed_experts(),
687
num_redundant_experts=getattr(self, "num_redundant_experts", 0),
688
)
689
return expert_params_mapping
vllm/model_executor/models/registry.py
+1
@@ -163,6 +163,7 @@ _TEXT_GENERATION_MODELS = {
163
"MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),
164
"NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
165
"NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
166
+ "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
167
"OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
168
"Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
169
"Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
vllm/transformers_utils/model_arch_config_convertor.py
+25
-2
@@ -81,6 +81,26 @@ class ModelArchConfigConvertorBase:
81
self.hf_text_config, attributes, default_factory=default_factory
82
)
83
84
+ def get_num_experts_from_block_configs(self) -> int:
85
+ """Check block_configs for heterogeneous models (e.g., NemotronH).
86
+
87
+ For heterogeneous models with varying expert counts per layer,
88
+ returns the MAX to ensure all expert weights can be loaded.
89
+ """
90
+ max_experts = 0
91
+ block_configs = getattr(self.hf_text_config, "block_configs", None)
92
+ if block_configs:
93
+ for block in block_configs:
94
+ if isinstance(block, dict):
95
+ if block.get("block_type", "") == "moe":
96
+ max_experts = max(max_experts, block.get("n_routed_experts", 0))
97
+ else:
98
+ if getattr(block, "block_type", "") == "moe":
99
+ max_experts = max(
100
+ max_experts, getattr(block, "n_routed_experts", 0)
101
+ )
102
+ return max_experts
103
+
104
def get_num_experts(self) -> int:
105
"""Returns the number of experts in the model."""
106
num_expert_names = [
@@ -89,13 +109,16 @@ class ModelArchConfigConvertorBase:
109
"n_routed_experts", # DeepSeek
110
"num_local_experts", # Mixtral
111
]
112
+
113
num_experts = getattr_iter(self.hf_text_config, num_expert_names, 0)
114
if isinstance(num_experts, list):
115
# Ernie VL's remote code uses list[int]...
116
# The values are always the same so we just take the first one.
117
return num_experts[0]
97
- # Coerce to 0 if explicitly set to None
98
- return num_experts or 0
118
+
119
+ if not num_experts:
120
+ num_experts = self.get_num_experts_from_block_configs()
121
+ return num_experts
122
123
@final
124
@classmethod