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, 02:30 UTC
[MoE] Enable Shared/Routed Overlap For Latent MoE (Nemotron-H) (#32790)
Signed-off-by: dafrimi <dafrimi@nvidia.com>
danielafrimi committed
Feb 2, 2026 at 16:18 UTC
0aca8b8c628e9a73ab8758d78c9c721bc703ee66
4 files changed
+303
-58
tests/kernels/moe/test_shared_fused_moe_routed_transform.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for SharedFusedMoE with routed_input_transform.
5
+
6
+Verifies that applying routed_input_transform inside SharedFusedMoE
7
+produces the same results as applying the transform manually outside.
8
+"""
9
+
10
+import pytest
11
+import torch
12
+import torch.nn as nn
13
+
14
+from vllm.config import VllmConfig, set_current_vllm_config
15
+from vllm.forward_context import set_forward_context
16
+from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
17
+
18
+
19
+class SimpleLinear(nn.Module):
20
+ """A simple linear transform mimicking latent projection in latent MoE."""
21
+
22
+ def __init__(self, in_features: int, out_features: int, dtype: torch.dtype):
23
+ super().__init__()
24
+ self.weight = nn.Parameter(
25
+ torch.randn(out_features, in_features, device="cuda", dtype=dtype) / 10
26
+ )
27
+
28
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
29
+ return nn.functional.linear(x, self.weight)
30
+
31
+
32
+class SimpleSharedExperts(nn.Module):
33
+ """A simple 2-layer MLP mimicking shared experts."""
34
+
35
+ def __init__(self, hidden_size: int, intermediate_size: int, dtype: torch.dtype):
36
+ super().__init__()
37
+ self.up = nn.Linear(
38
+ hidden_size, intermediate_size * 2, bias=False, device="cuda", dtype=dtype
39
+ )
40
+ self.down = nn.Linear(
41
+ intermediate_size, hidden_size, bias=False, device="cuda", dtype=dtype
42
+ )
43
+ with torch.no_grad():
44
+ self.up.weight.div_(10)
45
+ self.down.weight.div_(10)
46
+
47
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
48
+ gate_up = self.up(x)
49
+ gate, up = gate_up.chunk(2, dim=-1)
50
+ return self.down(nn.functional.silu(gate) * up)
51
+
52
+
53
+@pytest.fixture(autouse=True)
54
+def setup_cuda():
55
+ if not torch.cuda.is_available():
56
+ pytest.skip("CUDA not available")
57
+ torch.set_default_device("cuda")
58
+
59
+
60
+@pytest.mark.parametrize("num_tokens", [1, 32])
61
+@pytest.mark.parametrize("hidden_size,latent_size", [(256, 128), (128, 64)])
62
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
63
+def test_routed_input_transform_inside_vs_outside(
64
+ num_tokens: int,
65
+ hidden_size: int,
66
+ latent_size: int,
67
+ dtype: torch.dtype,
68
+ dist_init,
69
+ workspace_init,
70
+):
71
+ """Compare SharedFusedMoE with transform inside vs manually applying outside.
72
+ Method A (inside): SharedFusedMoE with routed_input_transform
73
+ Method B (outside): Manually transform, then SharedFusedMoE without transform
74
+ """
75
+ torch.manual_seed(42)
76
+
77
+ num_experts = 8
78
+ top_k = 2
79
+ intermediate_size = hidden_size * 2
80
+
81
+ vllm_config = VllmConfig()
82
+ vllm_config.compilation_config.static_forward_context = dict()
83
+
84
+ shared_experts = SimpleSharedExperts(hidden_size, intermediate_size, dtype)
85
+ routed_transform = SimpleLinear(hidden_size, latent_size, dtype)
86
+
87
+ with set_current_vllm_config(vllm_config):
88
+ # Method A: SharedFusedMoE WITH routed_input_transform
89
+ moe_with_transform = SharedFusedMoE(
90
+ shared_experts=shared_experts,
91
+ routed_input_transform=routed_transform,
92
+ num_experts=num_experts,
93
+ top_k=top_k,
94
+ hidden_size=latent_size,
95
+ intermediate_size=intermediate_size,
96
+ reduce_results=False,
97
+ renormalize=True,
98
+ params_dtype=dtype,
99
+ tp_size=1,
100
+ dp_size=1,
101
+ pcp_size=1,
102
+ prefix="moe_with_transform",
103
+ )
104
+
105
+ # Method B: SharedFusedMoE WITHOUT routed_input_transform
106
+ # Note: shared_experts=None because when transform is done outside,
107
+ moe_without_transform = SharedFusedMoE(
108
+ shared_experts=None,
109
+ routed_input_transform=None,
110
+ num_experts=num_experts,
111
+ top_k=top_k,
112
+ hidden_size=latent_size,
113
+ intermediate_size=intermediate_size,
114
+ reduce_results=False,
115
+ renormalize=True,
116
+ params_dtype=dtype,
117
+ tp_size=1,
118
+ dp_size=1,
119
+ pcp_size=1,
120
+ prefix="moe_without_transform",
121
+ )
122
+
123
+ with torch.no_grad():
124
+ moe_without_transform.w13_weight.copy_(moe_with_transform.w13_weight)
125
+ moe_without_transform.w2_weight.copy_(moe_with_transform.w2_weight)
126
+
127
+ moe_with_transform.quant_method.process_weights_after_loading(
128
+ moe_with_transform
129
+ )
130
+ moe_without_transform.quant_method.process_weights_after_loading(
131
+ moe_without_transform
132
+ )
133
+
134
+ hidden_states = torch.randn(num_tokens, hidden_size, device="cuda", dtype=dtype)
135
+ router_logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)
136
+
137
+ with set_forward_context(None, vllm_config, num_tokens=num_tokens):
138
+ shared_out_A, routed_out_A = moe_with_transform(
139
+ hidden_states, router_logits
140
+ )
141
+
142
+ transformed_hidden = routed_transform(hidden_states)
143
+ shared_out_B, routed_out_B = moe_without_transform(
144
+ transformed_hidden, router_logits
145
+ )
146
+
147
+ torch.testing.assert_close(
148
+ routed_out_A,
149
+ routed_out_B,
150
+ atol=1e-3,
151
+ rtol=1e-3,
152
+ msg="Routed output should match: transform inside vs outside",
153
+ )
154
+
155
+ expected_shared_out = shared_experts(hidden_states)
156
+
157
+ torch.testing.assert_close(
158
+ shared_out_A,
159
+ expected_shared_out,
160
+ atol=1e-3,
161
+ rtol=1e-3,
162
+ )
vllm/model_executor/layers/fused_moe/layer.py
+89
-16
@@ -1,8 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Callable, Iterable
5
-from contextlib import nullcontext
4
+from collections.abc import Callable, Generator, Iterable
5
+from contextlib import contextmanager, nullcontext
6
from enum import Enum
7
from typing import Literal, cast, get_args, overload
8
@@ -351,6 +351,10 @@ class FusedMoE(CustomOp):
351
"Enabled separate cuda stream for MoE shared_experts", scope="local"
352
)
353
354
+ # For latent MoE: stores original hidden_states before routed_input_transform
355
+ # so shared_experts can use it for cloning (they need original dimension)
356
+ self._shared_experts_input: torch.Tensor | None = None
357
+
358
if params_dtype is None:
359
params_dtype = torch.get_default_dtype()
360
self.params_dtype = params_dtype
@@ -664,6 +668,39 @@ class FusedMoE(CustomOp):
668
def gate(self) -> torch.nn.Module | None:
669
return None
670
671
+ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor:
672
+ """Hook to transform hidden_states before passing to routed experts.
673
+ For latent MoE: transforms [S, hidden_size] → [S, moe_latent_size].
674
+ The original hidden_states is saved in _shared_experts_input so
675
+ shared_experts still receive the original [S, hidden_size].
676
+
677
+ Override in subclasses (e.g., SharedFusedMoE) for latent MoE.
678
+ """
679
+ return hidden_states
680
+
681
+ @contextmanager
682
+ def _set_shared_experts_input(
683
+ self, value: torch.Tensor | None
684
+ ) -> Generator[None, None, None]:
685
+ """Context manager to safely set/clear _shared_experts_input."""
686
+ self._shared_experts_input = value
687
+ try:
688
+ yield
689
+ finally:
690
+ self._shared_experts_input = None
691
+
692
+ def _get_shared_experts_input(self, hidden_states: torch.Tensor) -> torch.Tensor:
693
+ """Get input for shared experts.
694
+
695
+ For latent MoE: shared_experts need original [S, hidden_size],
696
+ not the transformed [S, latent_size] used by routed experts.
697
+ """
698
+ return (
699
+ self._shared_experts_input
700
+ if self._shared_experts_input is not None
701
+ else hidden_states
702
+ )
703
+
704
@property
705
def tp_size(self):
706
return self.moe_parallel_config.tp_size
@@ -855,9 +892,11 @@ class FusedMoE(CustomOp):
892
if use_shared_experts_stream:
893
assert self.shared_experts_stream is not None
894
895
+ shared_experts_input = self._get_shared_experts_input(hidden_states)
896
+
897
# Clone BEFORE switching streams to avoid race condition
898
# where routed_expert kernel may mutate hidden_states.
860
- hidden_states_clone = hidden_states.clone()
899
+ hidden_states_clone = shared_experts_input.clone()
900
901
# Record that the clone will be used by shared_experts_stream
902
# to avoid gc issue from deallocation of hidden_states_clone
@@ -1537,11 +1576,20 @@ class FusedMoE(CustomOp):
1576
hidden_states: torch.Tensor,
1577
router_logits: torch.Tensor,
1578
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1540
- og_hidden_states = hidden_states.shape[-1]
1541
- if self.hidden_size != og_hidden_states:
1579
+ # For latent MoE: save ORIGINAL hidden_states before transform
1580
+ # (shared_experts need original dimension, routed experts use transformed)
1581
+ original_hidden_states = hidden_states
1582
+ original_hidden_dim = hidden_states.shape[-1]
1583
+
1584
+ # Apply transform for routed experts (e.g., latent projection for latent MoE)
1585
+ hidden_states = self.apply_routed_input_transform(hidden_states)
1586
+
1587
+ # This is the dimension after transform (for routed expert output slicing)
1588
+ transformed_hidden_dim = hidden_states.shape[-1]
1589
+ if self.hidden_size != transformed_hidden_dim:
1590
hidden_states = F.pad(
1591
hidden_states,
1544
- (0, self.hidden_size - og_hidden_states),
1592
+ (0, self.hidden_size - transformed_hidden_dim),
1593
mode="constant",
1594
value=0.0,
1595
)
@@ -1576,22 +1624,31 @@ class FusedMoE(CustomOp):
1624
fused_output = torch.ops.vllm.moe_forward(
1625
hidden_states, router_logits, encode_layer_name()
1626
)
1579
- return reduce_output(fused_output)[..., :og_hidden_states]
1627
+ return reduce_output(fused_output)[..., :transformed_hidden_dim]
1628
else:
1629
if current_platform.is_tpu() or current_platform.is_cpu():
1630
# TODO: Once the OOM issue for the TPU backend is resolved, we
1631
# will switch to using the moe_forward custom op.
1632
# Note: CPU doesn't require wrapped forward_impl.
1585
- shared_output, fused_output = self.forward_impl(
1586
- hidden_states, router_logits
1587
- )
1633
+ with self._set_shared_experts_input(original_hidden_states):
1634
+ shared_output, fused_output = self.forward_impl(
1635
+ hidden_states, router_logits
1636
+ )
1637
else:
1638
+ # Custom op handles setting/clearing _shared_experts_input internally
1639
+ # We pass original tensor for shared experts (not transformed)
1640
shared_output, fused_output = torch.ops.vllm.moe_forward_shared(
1590
- hidden_states, router_logits, encode_layer_name()
1641
+ hidden_states,
1642
+ router_logits,
1643
+ encode_layer_name(),
1644
+ original_hidden_states,
1645
)
1646
+
1647
+ # shared_output uses original dimension (before transform)
1648
+ # fused_output uses transformed dimension (after transform)
1649
return (
1593
- reduce_output(shared_output)[..., :og_hidden_states],
1594
- reduce_output(fused_output)[..., :og_hidden_states],
1650
+ reduce_output(shared_output)[..., :original_hidden_dim],
1651
+ reduce_output(fused_output)[..., :transformed_hidden_dim],
1652
)
1653
1654
@property
@@ -1831,7 +1888,8 @@ class FusedMoE(CustomOp):
1888
# because matrix multiply maybe modify the hidden_states.
1889
if has_separate_shared_experts and not use_shared_experts_stream:
1890
assert self.shared_experts is not None
1834
- shared_output = self.shared_experts(hidden_states)
1891
+ shared_input = self._get_shared_experts_input(hidden_states)
1892
+ shared_output = self.shared_experts(shared_input)
1893
1894
# NOTE: Similar with DP, PCP also needs dispatch and combine. For
1895
# simplicity, AgRsAll2All was added separately for PCP here. Maybe
@@ -2021,19 +2079,34 @@ def moe_forward_shared(
2079
hidden_states: torch.Tensor,
2080
router_logits: torch.Tensor,
2081
layer_name: str,
2082
+ shared_experts_input: torch.Tensor | None = None,
2083
) -> tuple[torch.Tensor, torch.Tensor]:
2084
self = get_layer_from_name(layer_name)
2085
assert self.shared_experts is not None
2027
- return self.forward_impl(hidden_states, router_logits)
2086
+
2087
+ # Set here because torch.compile skips forward_native() setup code
2088
+ # and calls this op directly. forward_impl() reads from this var.
2089
+ with self._set_shared_experts_input(shared_experts_input):
2090
+ return self.forward_impl(hidden_states, router_logits)
2091
2092
2093
def moe_forward_shared_fake(
2094
hidden_states: torch.Tensor,
2095
router_logits: torch.Tensor,
2096
layer_name: str,
2097
+ shared_experts_input: torch.Tensor | None = None,
2098
) -> tuple[torch.Tensor, torch.Tensor]:
2035
- shared_out = torch.empty_like(hidden_states)
2099
+ # Output shapes:
2100
+ # - fused_out: same as hidden_states (routed experts use transformed size)
2101
+ # - shared_out: same as shared_experts_input if provided, else same as hidden_states
2102
+ # (For latent MoE: shared experts use original hidden_size, not latent size)
2103
fused_out = torch.empty_like(hidden_states)
2104
+
2105
+ if shared_experts_input is not None:
2106
+ shared_out = torch.empty_like(shared_experts_input)
2107
+ else:
2108
+ shared_out = torch.empty_like(hidden_states)
2109
+
2110
return shared_out, fused_out
2111
2112
vllm/model_executor/layers/fused_moe/shared_fused_moe.py
+22
@@ -23,10 +23,12 @@ class SharedFusedMoE(FusedMoE):
23
shared_experts: torch.nn.Module | None,
24
gate: torch.nn.Module | None = None,
25
use_overlapped: bool = True,
26
+ routed_input_transform: torch.nn.Module | None = None,
27
**kwargs,
28
):
29
super().__init__(**kwargs)
30
self._shared_experts = shared_experts
31
+ self._routed_input_transform = routed_input_transform
32
33
# Disable shared expert overlap if:
34
# - we are using eplb with non-default backend, because of correctness issues
@@ -56,6 +58,26 @@ class SharedFusedMoE(FusedMoE):
58
def is_internal_router(self) -> bool:
59
return self.gate is not None
60
61
+ def apply_routed_input_transform(self, hidden_states: torch.Tensor) -> torch.Tensor:
62
+ """Apply transform for routed experts (e.g., latent projection).
63
+
64
+ This is called by FusedMoE.forward_native. The original hidden_states
65
+ is saved separately so shared experts get [S, hidden_size] while
66
+ routed experts get the transformed [S, moe_latent_size].
67
+
68
+ TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be
69
+ moved inside SharedFusedMoE to all-reduce on the smaller latent
70
+ dimension.
71
+ """
72
+ if self._routed_input_transform is not None:
73
+ result = self._routed_input_transform(hidden_states)
74
+ # ReplicatedLinear returns (output, extra_bias) tuple.
75
+ # We only need the output tensor; extra_bias is not used here.
76
+ if isinstance(result, tuple):
77
+ return result[0]
78
+ return result
79
+ return hidden_states
80
+
81
def forward(
82
self,
83
hidden_states: torch.Tensor,
vllm/model_executor/models/nemotron_h.py
+30
-42
@@ -188,10 +188,29 @@ class NemotronHMoE(nn.Module):
188
prefix=f"{prefix}.shared_experts",
189
)
190
191
+ if self.use_latent_moe:
192
+ self.fc1_latent_proj = ReplicatedLinear(
193
+ input_size=config.hidden_size,
194
+ output_size=self.moe_hidden_size,
195
+ bias=config.mlp_bias,
196
+ quant_config=quant_config,
197
+ disable_tp=self.is_sequence_parallel,
198
+ prefix=f"{prefix}.fc1_latent_proj",
199
+ )
200
+ self.fc2_latent_proj = ReplicatedLinear(
201
+ input_size=self.moe_hidden_size,
202
+ output_size=config.hidden_size,
203
+ bias=config.mlp_bias,
204
+ quant_config=quant_config,
205
+ disable_tp=self.is_sequence_parallel,
206
+ prefix=f"{prefix}.fc2_latent_proj",
207
+ )
208
+ else:
209
+ self.fc1_latent_proj = None
210
+ self.fc2_latent_proj = None
211
+
212
self.experts = SharedFusedMoE(
192
- # TODO: make it possible for shared experts to have
193
- # different input in SharedFusedMoE
194
- shared_experts=self.shared_experts if not self.use_latent_moe else None,
213
+ shared_experts=self.shared_experts,
214
num_experts=config.n_routed_experts,
215
top_k=config.num_experts_per_tok,
216
hidden_size=self.moe_hidden_size,
@@ -211,30 +230,9 @@ class NemotronHMoE(nn.Module):
230
num_redundant_experts=self.n_redundant_experts,
231
is_sequence_parallel=self.is_sequence_parallel,
232
router_logits_dtype=router_logits_dtype,
233
+ routed_input_transform=self.fc1_latent_proj,
234
)
235
216
- if self.use_latent_moe:
217
- self.fc1_latent_proj = ReplicatedLinear(
218
- input_size=config.hidden_size,
219
- output_size=self.moe_hidden_size,
220
- bias=config.mlp_bias,
221
- quant_config=quant_config,
222
- disable_tp=self.is_sequence_parallel,
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
-
236
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
237
num_tokens, hidden_dim = hidden_states.shape
238
hidden_states = hidden_states.view(-1, hidden_dim)
@@ -244,38 +242,28 @@ class NemotronHMoE(nn.Module):
242
243
# router_logits: (num_tokens, n_experts)
244
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)
245
253
- fused_moe_out = self.experts(
246
+ # SharedFusedMoE handles:
247
+ # - shared experts (with original hidden_states)
248
+ # - routed_input_transform (fc1_latent_proj) for latent MoE
249
+ # - multistream parallelism between shared and routed experts
250
+ shared_output, final_hidden_states = self.experts(
251
hidden_states=hidden_states, router_logits=router_logits
252
)
253
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
-
254
# Fix FP16 overflow
255
# See DeepseekV2DecoderLayer for more details.
256
if hidden_states.dtype != torch.float16:
257
final_hidden_states *= self.routed_scaling_factor
258
elif self.shared_experts is not None:
267
- assert shared_output is not None
259
shared_output *= 1.0 / self.routed_scaling_factor
260
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.
261
+ # TODO: See SharedFusedMoE.apply_routed_input_transform
262
+ # for bandwidth optimization
263
if self.use_latent_moe:
264
final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states)
265
266
if self.shared_experts is not None:
278
- assert shared_output is not None
267
final_hidden_states += shared_output
268
269
if self.is_sequence_parallel: