@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fcd5306f65

Fork 0

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 11, 21:15 UTC

Add latent MoE support (#30203)

Signed-off-by: Shahar Mor <smor@nvidia.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

shaharmor98 committed Dec 8, 2025 at 19:35 UTC fcd5306f65876f72122d7e5852b6000738498d7e

2 files changed +58 -6

vllm/model_executor/models/nemotron_h.py

+56 -6

@@ -83,6 +83,7 @@ class NemotronHMLP(nn.Module):

83 def __init__(

84 self,

85 config: NemotronHConfig,

86 + hidden_size: int,

87 intermediate_size: int,

88 quant_config: QuantizationConfig | None = None,

89 bias: bool = False,

@@ -93,7 +94,7 @@ class NemotronHMLP(nn.Module):

94 super().__init__()

96 self.up_proj = ColumnParallelLinear(

96 - input_size=config.hidden_size,

97 + input_size=hidden_size,

98 output_size=intermediate_size,

99 bias=bias,

100 quant_config=quant_config,

@@ -102,7 +103,7 @@ class NemotronHMLP(nn.Module):

103 )

104 self.down_proj = RowParallelLinear(

105 input_size=intermediate_size,

105 - output_size=config.hidden_size,

106 + output_size=hidden_size,

107 bias=bias,

108 quant_config=quant_config,

109 reduce_results=reduce_results,

@@ -135,6 +136,10 @@ class NemotronHMoE(nn.Module):

136 self.ep_size = self.ep_group.size()

137 self.n_routed_experts: int = config.n_routed_experts

138 self.n_shared_experts: int = config.n_shared_experts

139 + self.use_latent_moe: bool = getattr(config, "moe_latent_size", None) is not None

140 + self.moe_hidden_size: int = (

141 + config.moe_latent_size if self.use_latent_moe else config.hidden_size

142 + )

143

144 self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe

145

@@ -172,6 +177,7 @@ class NemotronHMoE(nn.Module):

177

178 self.shared_experts = NemotronHMLP(

179 config=config,

180 + hidden_size=config.hidden_size,

181 intermediate_size=intermediate_size,

182 quant_config=quant_config,

183 reduce_results=False,

@@ -180,10 +186,12 @@ class NemotronHMoE(nn.Module):

186 )

187

188 self.experts = SharedFusedMoE(

183 - shared_experts=self.shared_experts,

189 + # TODO: make it possible for shared experts to have

190 + # different input in SharedFusedMoE

191 + shared_experts=self.shared_experts if not self.use_latent_moe else None,

192 num_experts=config.n_routed_experts,

193 top_k=config.num_experts_per_tok,

186 - hidden_size=config.hidden_size,

194 + hidden_size=self.moe_hidden_size,

195 intermediate_size=config.moe_intermediate_size,

196 reduce_results=False,

197 renormalize=config.norm_topk_prob,

@@ -201,6 +209,32 @@ class NemotronHMoE(nn.Module):

209 is_sequence_parallel=self.is_sequence_parallel,

210 )

211

212 + if self.use_latent_moe:

213 + # TODO: check if using ReplicatedLinear is better than

214 + # ColumnParallelLinear + all_gather

215 + self.fc1_latent_proj = ColumnParallelLinear(

216 + input_size=config.hidden_size,

217 + output_size=self.moe_hidden_size,

218 + bias=config.mlp_bias,

219 + quant_config=quant_config,

220 + disable_tp=self.is_sequence_parallel,

221 + # We need to gather the output to prepare input for moe

222 + gather_output=True,

223 + prefix=f"{prefix}.fc1_latent_proj",

224 + )

225 + self.fc2_latent_proj = ReplicatedLinear(

226 + input_size=self.moe_hidden_size,

227 + output_size=config.hidden_size,

228 + bias=config.mlp_bias,

229 + quant_config=quant_config,

230 + disable_tp=self.is_sequence_parallel,

231 + prefix=f"{prefix}.fc2_latent_proj",

232 + )

233 +

234 + else:

235 + self.fc1_latent_proj = None

236 + self.fc2_latent_proj = None

237 +

238 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

239 num_tokens, hidden_dim = hidden_states.shape

240 hidden_states = hidden_states.view(-1, hidden_dim)

@@ -210,12 +244,20 @@ class NemotronHMoE(nn.Module):

244

245 # router_logits: (num_tokens, n_experts)

246 router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))

247 + shared_output = None

248 + if self.use_latent_moe:

249 + if self.shared_experts is not None:

250 + shared_output = self.shared_experts(hidden_states)

251 + hidden_states, _ = self.fc1_latent_proj(hidden_states)

252

253 fused_moe_out = self.experts(

254 hidden_states=hidden_states, router_logits=router_logits

255 )

256

218 - shared_output, final_hidden_states = fused_moe_out

257 + if self.use_latent_moe:

258 + _, final_hidden_states = fused_moe_out

259 + else:

260 + shared_output, final_hidden_states = fused_moe_out

261

262 # Fix FP16 overflow

263 # See DeepseekV2DecoderLayer for more details.

@@ -225,6 +267,13 @@ class NemotronHMoE(nn.Module):

267 assert shared_output is not None

268 shared_output *= 1.0 / self.routed_scaling_factor

269

270 + # TODO: currently latent up_proj is done before all-reduce for simplicity.

271 + # if and when shared experts will be part of SharedFusedMoE,

272 + # we should do the up_proj after all-reduce,

273 + # to have the all-reduce in the smaller latent dimension.

274 + if self.use_latent_moe:

275 + final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states)

276 +

277 if self.shared_experts is not None:

278 assert shared_output is not None

279 final_hidden_states += shared_output

@@ -268,6 +317,7 @@ class NemotronHMLPDecoderLayer(nn.Module):

317

318 self.mixer = NemotronHMLP(

319 config,

320 + hidden_size=config.hidden_size,

321 intermediate_size=intermediate_size,

322 quant_config=quant_config,

323 bias=config.mlp_bias,

@@ -846,5 +896,5 @@ class NemotronHForCausalLM(

896 return logits

897

898 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

849 - loader = AutoWeightsLoader(self)

899 + loader = AutoWeightsLoader(self, skip_prefixes=["mtp"])

900 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/transformers_utils/configs/nemotron_h.py

@@ -189,6 +189,7 @@ class NemotronHConfig(PretrainedConfig):

189 n_shared_experts=1,

190 moe_intermediate_size=7688,

191 moe_shared_expert_intermediate_size=7688,

192 + moe_latent_size=None,

193 num_experts_per_tok=2,

194 routed_scaling_factor=1.0,

195 n_group=1,

@@ -254,6 +255,7 @@ class NemotronHConfig(PretrainedConfig):

255 self.n_shared_experts = n_shared_experts

256 self.moe_intermediate_size = moe_intermediate_size

257 self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size # noqa: E501

258 + self.moe_latent_size = moe_latent_size

259 self.num_experts_per_tok = num_experts_per_tok

260 self.routed_scaling_factor = routed_scaling_factor

261 self.n_group = n_group