Add MoE output contract for MoE tail fusion (#51407)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Aug 11, 2026 at 22:23 UTC
dd7cc85f1720ebcb65b1822cbc5054af66ae292d
4 files changed
+127
-13
vllm/model_executor/layers/fused_moe/config.py
+19
@@ -1311,6 +1311,13 @@ class FusedMoEConfig:
1311
# Only honored on the non-reduced (late-AR) TP path. Default False.
1312
skip_final_all_reduce: bool = False
1313
1314
+ # When True, experts that can stop after GEMM2 are allowed to hand back an
1315
+ # UnfinalizedMoEOutput instead of finalized states, leaving the top-k
1316
+ # reduction to fuse into the consumer. Set by layers that have such a
1317
+ # consumer; read through `use_deferred_moe_finalize`, which applies the
1318
+ # guards. Kernels without the capability ignore it. Default False.
1319
+ defer_moe_finalize: bool = False
1320
+
1321
# SwiGLU clamp limit. When set, backends that do not implement the clamp
1322
# are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
1323
# cannot silently select one and drop the clamp.
@@ -1428,6 +1435,18 @@ class FusedMoEConfig:
1435
def use_ep(self):
1436
return self.moe_parallel_config.use_ep
1437
1438
+ @property
1439
+ def use_deferred_moe_finalize(self) -> bool:
1440
+ """Whether experts may return an unfinalized output on this deployment.
1441
+
1442
+ Evaluated on read rather than in ``__post_init__`` because
1443
+ ``defer_moe_finalize`` is set after construction, like
1444
+ ``skip_final_all_reduce``.
1445
+ """
1446
+ # Without TP there is no all-reduce for the top-k reduction to fuse
1447
+ # into, so the deferred form buys nothing.
1448
+ return self.defer_moe_finalize and self.tp_size > 1
1449
+
1450
@property
1451
def use_deepep_ht_kernels(self):
1452
return self.moe_parallel_config.use_deepep_ht_kernels
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+25
-6
@@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
RoutingMethodType,
13
)
14
+from vllm.model_executor.layers.fused_moe.moe_output import (
15
+ UnfinalizedMoEOutput,
16
+)
17
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
18
TopKWeightAndReduceNoOP,
19
)
@@ -266,13 +269,18 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
269
e_score_correction_bias: torch.Tensor | None = None,
270
routed_scaling_factor: float | None = None,
271
topk_group: int | None = None,
269
- ) -> torch.Tensor:
272
+ ) -> torch.Tensor | UnfinalizedMoEOutput:
273
import flashinfer
274
275
assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
276
277
+ num_tokens = hidden_states.shape[0]
278
+ # The runner divides by the token count on the host, so an idle rank's
279
+ # dummy 0-token forward has to keep the finalized (empty) form.
280
+ defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
281
+
282
routing_replay_out = self._maybe_make_routing_replay_buffer(
275
- num_tokens=hidden_states.shape[0],
283
+ num_tokens=num_tokens,
284
device=hidden_states.device,
285
)
286
out = flashinfer.fused_moe.trtllm_bf16_moe(
@@ -293,8 +301,19 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
301
activation_type=activation_to_flashinfer_int(activation),
302
tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
303
routing_replay_out=routing_replay_out,
304
+ do_finalize=not defer,
305
)
297
- self._maybe_dispatch_routing_replay(
298
- routing_replay_out, num_tokens=hidden_states.shape[0]
299
- )
300
- return out
306
+ self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
307
+ if defer:
308
+ # flashinfer returns a flat permute map; the protocol wants
309
+ # [num_tokens, top_k] so consumers can read top_k from its shape.
310
+ return UnfinalizedMoEOutput(
311
+ gemm2_permuted=out[0],
312
+ expert_weights=out[1],
313
+ expanded_idx_to_permuted_idx=out[2]
314
+ .to(torch.int32)
315
+ .view(num_tokens, self.topk),
316
+ )
317
+ # do_finalize=True yields the finalized states (a bare tensor on some
318
+ # FlashInfer versions, a single-element list on others).
319
+ return out[0] if isinstance(out, (list, tuple)) else out
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+23
-7
@@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import (
13
FusedMoEQuantConfig,
14
RoutingMethodType,
15
)
16
+from vllm.model_executor.layers.fused_moe.moe_output import (
17
+ UnfinalizedMoEOutput,
18
+)
19
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
20
TopKWeightAndReduceNoOP,
21
)
@@ -518,7 +521,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
521
e_score_correction_bias: torch.Tensor | None = None,
522
routed_scaling_factor: float | None = None,
523
topk_group: int | None = None,
521
- ) -> torch.Tensor:
524
+ ) -> torch.Tensor | UnfinalizedMoEOutput:
525
import flashinfer
526
527
assert self._supports_activation(activation)
@@ -543,8 +546,13 @@ class TrtLlmNvFp4ExpertsMonolithic(
546
547
output1_scale_gate_scalar = self.quant_config.g1_alphas
548
549
+ num_tokens = hidden_states.shape[0]
550
+ # The runner divides by the token count on the host, so an idle rank's
551
+ # dummy 0-token forward has to keep the finalized (empty) form.
552
+ defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
553
+
554
routing_replay_out = self._maybe_make_routing_replay_buffer(
547
- num_tokens=hidden_states.shape[0],
555
+ num_tokens=num_tokens,
556
device=hidden_states.device,
557
)
558
# Invoke kernel.
@@ -578,13 +586,21 @@ class TrtLlmNvFp4ExpertsMonolithic(
586
local_num_experts=self.local_num_experts,
587
routed_scaling_factor=routed_scaling_factor,
588
routing_method_type=self.routing_method_type,
581
- do_finalize=True,
589
+ do_finalize=not defer,
590
activation_type=activation_to_flashinfer_int(activation),
591
per_token_scale=per_token_scale,
592
tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
593
routing_replay_out=routing_replay_out,
586
- )[0]
587
- self._maybe_dispatch_routing_replay(
588
- routing_replay_out, num_tokens=hidden_states.shape[0]
594
)
590
- return result
595
+ self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
596
+ if defer:
597
+ # flashinfer returns a flat permute map; the protocol wants
598
+ # [num_tokens, top_k] so consumers can read top_k from its shape.
599
+ return UnfinalizedMoEOutput(
600
+ gemm2_permuted=result[0],
601
+ expert_weights=result[1],
602
+ expanded_idx_to_permuted_idx=result[2]
603
+ .to(torch.int32)
604
+ .view(num_tokens, self.topk),
605
+ )
606
+ return result[0]
vllm/model_executor/layers/fused_moe/moe_output.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Output contract between a MoE layer and a consumer that fuses its tail."""
4
+
5
+from dataclasses import dataclass
6
+
7
+import torch
8
+
9
+
10
+@dataclass
11
+class UnfinalizedMoEOutput:
12
+ """Unfinalized output of a monolithic MoE kernel.
13
+
14
+ Kernels that can stop after GEMM2 (the TRTLLM-Gen ``do_finalize=False`` path)
15
+ hand back their permuted, unweighted output plus the routing weights and the
16
+ permute map, so that the top-k reduction can be fused with whatever follows
17
+ -- the shared-expert add and the tensor-parallel all-reduce -- instead of
18
+ running as its own kernel.
19
+
20
+ The buffers are consumed as-is by the fused kernels, which index
21
+ ``gemm2_permuted`` by row: it must be densely packed at ``hidden_dim``, and
22
+ its row count (an autotuner-dependent padded value) is never referenced.
23
+ """
24
+
25
+ # [num_permuted_rows, hidden_dim] permuted, unweighted GEMM2 output.
26
+ gemm2_permuted: torch.Tensor
27
+ # [num_tokens, top_k] routing weights, in the activation dtype -- consumers
28
+ # index them as such, and a buffer typed wider than its contents reads as
29
+ # garbage. These already carry routed_scaling_factor for routing methods that
30
+ # fold it in.
31
+ expert_weights: torch.Tensor
32
+ # [num_tokens, top_k] int32 permute map; -1 = expert not local to this rank.
33
+ expanded_idx_to_permuted_idx: torch.Tensor
34
+
35
+
36
+@dataclass
37
+class MoEOutput:
38
+ """A MoE layer's output with its final reduction still open.
39
+
40
+ Returned by layers whose MoE runs un-reduced (``reduce_results=False``) so
41
+ that the consumer -- typically the next layer's RMSNorm -- can fuse the
42
+ tensor-parallel all-reduce into itself instead of paying for a standalone
43
+ one. Keeping the shared-expert output and the routed scale separate leaves
44
+ that reduction the consumer's to schedule; when the routed output is still
45
+ unfinalized, the top-k reduction is open too and can fold into the same
46
+ kernel.
47
+
48
+ Producers only leave the routed output unfinalized when a fused consumer can
49
+ actually take that form -- the token ceiling and topology support are theirs
50
+ to check -- so an ``UnfinalizedMoEOutput`` here means the fused path applies,
51
+ and a consumer need not re-derive that.
52
+ """
53
+
54
+ # Un-reduced routed output, either finalized or not.
55
+ routed: torch.Tensor | UnfinalizedMoEOutput
56
+ # Un-reduced shared-expert output, folded in by whoever finalizes.
57
+ shared_output: torch.Tensor | None = None
58
+ # Applied to the routed sum before the shared add. Already folded into the
59
+ # routing weights for methods that do so, in which case this is 1.0.
60
+ routed_scaling_factor: float = 1.0