[NemotronH] Add latent MoE support (#16227)

Signed-off-by: Roi Koren <roik@nvidia.com>

python/sglang/srt/models/nemotron_h.py

@@ -138,6 +138,10 @@ class NemotronHMoE(nn.Module):
         self.ep_size = self.ep_group.size()
         self.n_routed_experts = config.n_routed_experts
         self.n_shared_experts = config.n_shared_experts
+        self.use_latent_moe = getattr(config, "moe_latent_size", None) is not None
+        self.moe_hidden_size = (
+            config.moe_latent_size if self.use_latent_moe else config.hidden_size
+        )
 
         self.gate = ReplicatedLinear(
             config.hidden_size,
@@ -165,7 +169,7 @@ class NemotronHMoE(nn.Module):
             num_experts=config.n_routed_experts
             + get_global_server_args().ep_num_redundant_experts,
             top_k=config.num_experts_per_tok,
-            hidden_size=config.hidden_size,
+            hidden_size=self.moe_hidden_size,
             intermediate_size=config.moe_intermediate_size,
             reduce_results=False,
             quant_config=quant_config,
@@ -186,6 +190,26 @@ class NemotronHMoE(nn.Module):
         else:
             self.shared_experts = None
 
+        if self.use_latent_moe:
+            self.fc1_latent_proj = ColumnParallelLinear(
+                input_size=config.hidden_size,
+                output_size=self.moe_hidden_size,
+                bias=config.mlp_bias,
+                quant_config=quant_config,
+                gather_output=True,
+                prefix=f"{prefix}.fc1_latent_proj",
+            )
+            self.fc2_latent_proj = ReplicatedLinear(
+                input_size=self.moe_hidden_size,
+                output_size=config.hidden_size,
+                bias=config.mlp_bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.fc2_latent_proj",
+            )
+        else:
+            self.fc1_latent_proj = None
+            self.fc2_latent_proj = None
+
     def _forward_core(
         self,
         hidden_states: torch.Tensor,
@@ -206,6 +230,8 @@ class NemotronHMoE(nn.Module):
         else:
             shared_output = None
         topk_output = self.topk(hidden_states, router_logits)
+        if self.use_latent_moe:
+            hidden_states, _ = self.fc1_latent_proj(hidden_states)
         final_hidden_states = self.experts(hidden_states, topk_output)
         return final_hidden_states, shared_output
 
@@ -226,6 +252,8 @@ class NemotronHMoE(nn.Module):
             # router_scores: [num_tokens, num_experts]
             router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
             topk_output = self.topk(hidden_states, router_logits)
+            if self.use_latent_moe:
+                hidden_states, _ = self.fc1_latent_proj(hidden_states)
             final_hidden_states = self.experts(hidden_states, topk_output)
         get_current_device_stream_fast().wait_stream(alt_stream)
 
@@ -242,6 +270,9 @@ class NemotronHMoE(nn.Module):
             assert shared_output is not None
             shared_output *= 1.0 / self.routed_scaling_factor
 
+        if self.use_latent_moe:
+            final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states)
+
         if shared_output is not None:
             final_hidden_states += shared_output