@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7a2fdbaac4

[K3 Perf] Fuse MXFP4 top-k finalization into latent-tail, ~5% E2E latency reduction (#53152)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Wentao Ye committed Aug 21, 2026 at 13:29 UTC 7a2fdbaac44921f57b04e53d6f49ed594c2209b6
17 files changed +720 -139
benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py
+89 -17
@@ -21,6 +21,10 @@ Examples:
21 torchrun --nproc-per-node=8 \
22 benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail
23
24 + torchrun --nproc-per-node=8 \
25 + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail \
26 + --backend fused --routed-input both
27 +
28 For multi-node runs, launch one ``torchrun`` agent per node and use a shared
29 rendezvous endpoint.
30 """
@@ -50,6 +54,7 @@ from vllm.distributed.parallel_state import (
54 initialize_model_parallel,
55 set_custom_all_reduce,
56 )
57 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
58 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
59 from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail
60 from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import (
@@ -165,6 +170,18 @@ def parse_args() -> argparse.Namespace:
170 nargs="+",
171 default=[1, 5, 8, 16],
172 )
173 + whole_tail.add_argument(
174 + "--routed-input",
175 + choices=("finalized", "deferred", "both"),
176 + default="finalized",
177 + help="Select whether the fused tail performs the top-k finalize step.",
178 + )
179 + whole_tail.add_argument(
180 + "--top-k",
181 + type=int,
182 + default=8,
183 + help="Expert count used by the synthetic deferred-finalize input.",
184 + )
185 whole_tail.add_argument("--warmup-replays", type=int, default=20)
186 whole_tail.add_argument("--samples", type=int, default=51)
187 whole_tail.add_argument(
@@ -586,6 +603,36 @@ def make_inputs(
603 return routed, shared
604
605
606 +def make_deferred_input(
607 + routed: torch.Tensor,
608 + top_k: int,
609 +) -> UnfinalizedMoEOutput:
610 + num_tokens = routed.shape[0]
611 + gemm2_permuted = (
612 + routed[:, None, :]
613 + .expand(num_tokens, top_k, LATENT_SIZE)
614 + .contiguous()
615 + .view(num_tokens * top_k, LATENT_SIZE)
616 + )
617 + expert_weights = torch.zeros(
618 + num_tokens,
619 + top_k,
620 + dtype=routed.dtype,
621 + device=routed.device,
622 + )
623 + expert_weights[:, 0] = 1
624 + expanded_idx = torch.arange(
625 + num_tokens * top_k,
626 + dtype=torch.int32,
627 + device=routed.device,
628 + ).view(num_tokens, top_k)
629 + return UnfinalizedMoEOutput(
630 + gemm2_permuted=gemm2_permuted,
631 + expert_weights=expert_weights,
632 + expanded_idx_to_permuted_idx=expanded_idx,
633 + )
634 +
635 +
636 def make_reference(
637 routed: torch.Tensor,
638 shared: torch.Tensor,
@@ -628,6 +675,8 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None:
675 raise ValueError("--num-tokens values must be in [1, 16]")
676 if args.warmup_replays < 0 or args.samples <= 0:
677 raise ValueError("warmup replays must be nonnegative and samples positive")
678 + if args.top_k <= 0:
679 + raise ValueError("--top-k must be positive")
680 if args.skinny_max_num_tokens is not None and any(
681 not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens
682 ):
@@ -672,6 +721,11 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None:
721
722 use_reference = args.backend in ("reference", "both")
723 use_fused = args.backend in ("fused", "both")
724 + routed_input_modes = (
725 + ("finalized", "deferred")
726 + if args.routed_input == "both"
727 + else (args.routed_input,)
728 + )
729 fused_ops = []
730 if use_fused:
731 production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m
@@ -690,23 +744,33 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None:
744 for cutoff in cutoffs:
745 latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff
746 latent_moe_tail.KimiK3LatentMoETailOp._instances.clear()
693 - fused_ops.append(
694 - (
695 - cutoff,
696 - latent_moe_tail.KimiK3LatentMoETailOp.initialize(
697 - hidden_size=HIDDEN_SIZE,
698 - latent_size=LATENT_SIZE,
699 - dtype=torch.bfloat16,
700 - device=device,
701 - rms_eps=RMS_EPS,
702 - ),
747 + for routed_input_mode in routed_input_modes:
748 + fused_ops.append(
749 + (
750 + cutoff,
751 + routed_input_mode,
752 + latent_moe_tail.KimiK3LatentMoETailOp.initialize(
753 + hidden_size=HIDDEN_SIZE,
754 + latent_size=LATENT_SIZE,
755 + dtype=torch.bfloat16,
756 + device=device,
757 + rms_eps=RMS_EPS,
758 + experts_per_token=(
759 + args.top_k if routed_input_mode == "deferred" else 0
760 + ),
761 + ),
762 + )
763 )
704 - )
764 cutedsl_warmup()
765
766 results = []
767 for num_tokens in args.num_tokens:
768 routed, shared = make_inputs(num_tokens, rank, device)
769 + routed_inputs: dict[str, torch.Tensor | UnfinalizedMoEOutput] = {
770 + "finalized": routed
771 + }
772 + if "deferred" in routed_input_modes:
773 + routed_inputs["deferred"] = make_deferred_input(routed, args.top_k)
774 reference = make_reference(
775 routed,
776 shared,
@@ -724,17 +788,23 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None:
788 device_group=device_group,
789 cpu_group=cpu_group,
790 )
727 - for cutoff, fused_op in fused_ops:
791 + for cutoff, routed_input_mode, fused_op in fused_ops:
792 + routed_input = routed_inputs[routed_input_mode]
793
794 def fused(
730 - routed: torch.Tensor = routed,
795 + routed_input: torch.Tensor | UnfinalizedMoEOutput = routed_input,
796 shared: torch.Tensor = shared,
797 fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op,
798 ) -> torch.Tensor:
734 - return fused_op(routed, shared, rms_weight, up_weight)
799 + return fused_op(routed_input, shared, rms_weight, up_weight)
800
801 fused_graph, fused_output = capture_tail_graph(fused, cpu_group)
737 - fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}"
802 + fused_key_parts = ["fused"]
803 + if routed_input_mode != "finalized" or len(routed_input_modes) > 1:
804 + fused_key_parts.append(routed_input_mode)
805 + if len(cutoffs) > 1:
806 + fused_key_parts.extend(("skinny", "max", str(cutoff)))
807 + fused_key = "_".join(fused_key_parts)
808 result[fused_key] = benchmark_tail_graph(
809 fused_graph,
810 warmup_replays=args.warmup_replays,
@@ -747,7 +817,7 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None:
817 speedup = (
818 result["reference"]["median_us"] / result[fused_key]["median_us"]
819 )
750 - if len(fused_ops) == 1:
820 + if fused_key == "fused":
821 result["speedup"] = speedup
822 else:
823 result[f"{fused_key}_speedup"] = speedup
@@ -763,7 +833,9 @@ def benchmark_whole_tail(args: argparse.Namespace) -> None:
833 "cuda_version": torch.version.cuda,
834 "warmup_replays": args.warmup_replays,
835 "samples": args.samples,
766 - "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops],
836 + "routed_input": args.routed_input,
837 + "top_k": args.top_k,
838 + "skinny_max_num_tokens": sorted({cutoff for cutoff, _, _ in fused_ops}),
839 "skinny_configs": {
840 str(num_tokens): asdict(config)
841 for num_tokens, config in skinny_configs.items()
tests/models/kimi_k3/test_latent_moe_tail.py
+152
@@ -13,6 +13,7 @@ from tests.utils import (
13 multi_process_parallel,
14 )
15 from vllm.distributed import get_tp_group
16 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
17 from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
18 from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp
19 from vllm.platforms import current_platform
@@ -20,6 +21,57 @@ from vllm.platforms import current_platform
21 HIDDEN_SIZE = 7168
22 LATENT_SIZE = 3584
23 EPS = 0.1
24 +TOP_K = 8
25 +
26 +
27 +def _make_deferred_routed_output(
28 + num_tokens: int,
29 + device: torch.device,
30 +) -> tuple[UnfinalizedMoEOutput, torch.Tensor]:
31 + num_routes = num_tokens * TOP_K
32 + num_permuted_rows = num_routes + 7
33 + expanded_output = torch.randn(
34 + num_routes,
35 + LATENT_SIZE,
36 + device=device,
37 + dtype=torch.bfloat16,
38 + ).mul_(0.01)
39 + expert_weights = torch.rand(
40 + num_tokens,
41 + TOP_K,
42 + device=device,
43 + dtype=torch.bfloat16,
44 + )
45 + expert_weights.div_(expert_weights.sum(dim=-1, keepdim=True))
46 + expanded_idx = torch.randperm(num_permuted_rows, device=device)[:num_routes]
47 + gemm2_permuted = torch.empty(
48 + num_permuted_rows,
49 + LATENT_SIZE,
50 + device=device,
51 + dtype=torch.bfloat16,
52 + )
53 + gemm2_permuted[expanded_idx] = expanded_output
54 +
55 + finalized = torch.zeros(
56 + num_tokens,
57 + LATENT_SIZE,
58 + device=device,
59 + dtype=torch.bfloat16,
60 + )
61 + expanded_output = expanded_output.view(num_tokens, TOP_K, LATENT_SIZE)
62 + for slot in range(TOP_K):
63 + finalized.add_(expanded_output[:, slot] * expert_weights[:, slot, None])
64 +
65 + return (
66 + UnfinalizedMoEOutput(
67 + gemm2_permuted=gemm2_permuted,
68 + expert_weights=expert_weights,
69 + expanded_idx_to_permuted_idx=expanded_idx.to(torch.int32).view(
70 + num_tokens, TOP_K
71 + ),
72 + ),
73 + finalized,
74 + )
75
76
77 @ray.remote(num_gpus=1, max_calls=1)
@@ -131,6 +183,99 @@ def _run_latent_moe_tail_test(
183 )
184
185
186 +@ray.remote(num_gpus=1, max_calls=1)
187 +def _test_deferred_finalize_parity_worker(
188 + monkeypatch: pytest.MonkeyPatch,
189 + tp_size: int,
190 + pp_size: int,
191 + rank: int,
192 + distributed_init_port: str,
193 +) -> None:
194 + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
195 + device = torch.device(f"cuda:{rank}")
196 + torch.accelerator.set_device_index(device)
197 + init_test_distributed_environment(
198 + tp_size,
199 + pp_size,
200 + rank,
201 + distributed_init_port,
202 + )
203 +
204 + torch.manual_seed(1000 + rank)
205 + rms_weight = 1 + 0.1 * torch.randn(
206 + LATENT_SIZE,
207 + device=device,
208 + dtype=torch.bfloat16,
209 + )
210 + up_weight = (
211 + torch.randn(
212 + HIDDEN_SIZE,
213 + LATENT_SIZE,
214 + device=device,
215 + dtype=torch.bfloat16,
216 + )
217 + / LATENT_SIZE**0.5
218 + )
219 + finalized_op = KimiK3LatentMoETailOp.initialize(
220 + hidden_size=HIDDEN_SIZE,
221 + latent_size=LATENT_SIZE,
222 + dtype=torch.bfloat16,
223 + device=device,
224 + rms_eps=EPS,
225 + )
226 + deferred_op = KimiK3LatentMoETailOp.initialize(
227 + hidden_size=HIDDEN_SIZE,
228 + latent_size=LATENT_SIZE,
229 + dtype=torch.bfloat16,
230 + device=device,
231 + rms_eps=EPS,
232 + experts_per_token=TOP_K,
233 + )
234 + cutedsl_warmup()
235 +
236 + for iteration, num_tokens in enumerate((1, 5, 16)):
237 + torch.manual_seed(2000 + 100 * iteration + rank)
238 + deferred_output, finalized_output = _make_deferred_routed_output(
239 + num_tokens,
240 + device,
241 + )
242 + shared_output = torch.randn(
243 + num_tokens,
244 + HIDDEN_SIZE,
245 + device=device,
246 + dtype=torch.bfloat16,
247 + )
248 +
249 + expected = finalized_op(
250 + finalized_output,
251 + shared_output,
252 + rms_weight,
253 + up_weight,
254 + )
255 + actual = deferred_op(
256 + deferred_output,
257 + shared_output,
258 + rms_weight,
259 + up_weight,
260 + )
261 +
262 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
263 +
264 +
265 +def _run_deferred_finalize_parity_test(
266 + monkeypatch: pytest.MonkeyPatch,
267 + tp_size: int,
268 +) -> None:
269 + if not current_platform.is_device_capability_family(100):
270 + pytest.skip("K3 latent-MoE tail fusion requires SM100")
271 + multi_process_parallel(
272 + monkeypatch,
273 + tp_size,
274 + 1,
275 + _test_deferred_finalize_parity_worker,
276 + )
277 +
278 +
279 @multi_gpu_test(num_gpus=8)
280 def test_latent_moe_tail_tp8_matches_native_path(
281 monkeypatch: pytest.MonkeyPatch,
@@ -143,3 +288,10 @@ def test_latent_moe_tail_tp16_matches_native_path(
288 monkeypatch: pytest.MonkeyPatch,
289 ) -> None:
290 _run_latent_moe_tail_test(monkeypatch, 16)
291 +
292 +
293 +@multi_gpu_test(num_gpus=8)
294 +def test_latent_moe_tail_deferred_finalize_matches_finalized(
295 + monkeypatch: pytest.MonkeyPatch,
296 +) -> None:
297 + _run_deferred_finalize_parity_test(monkeypatch, 8)
vllm/model_executor/layers/fused_moe/config.py
+21 -3
@@ -1321,6 +1321,8 @@ class FusedMoEConfig:
1321 # consumer; read through `use_deferred_moe_finalize`, which applies the
1322 # guards. Kernels without the capability ignore it. Default False.
1323 defer_moe_finalize: bool = False
1324 + # Optional consumer capacity for deferred finalize. Negative means unbounded.
1325 + defer_moe_finalize_max_num_tokens: int = -1
1326
1327 # SwiGLU clamp limit. When set, backends that do not implement the clamp
1328 # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
@@ -1447,9 +1449,25 @@ class FusedMoEConfig:
1449 ``defer_moe_finalize`` is set after construction, like
1450 ``skip_final_all_reduce``.
1451 """
1450 - # Without TP there is no all-reduce for the top-k reduction to fuse
1451 - # into, so the deferred form buys nothing.
1452 - return self.defer_moe_finalize and self.tp_size > 1
1452 + # The consumer fuses a TP all-reduce. Other parallel modes require a
1453 + # combine or reduce-scatter after the experts and cannot defer it.
1454 + return (
1455 + self.defer_moe_finalize
1456 + and self.tp_size > 1
1457 + and self.dp_size == 1
1458 + and self.ep_size == 1
1459 + and self.pcp_size == 1
1460 + and not self.is_sequence_parallel
1461 + )
1462 +
1463 + def should_defer_moe_finalize(self, num_tokens: int) -> bool:
1464 + """Return whether this invocation may defer the top-k reduction."""
1465 + max_num_tokens = self.defer_moe_finalize_max_num_tokens
1466 + return (
1467 + self.use_deferred_moe_finalize
1468 + and num_tokens > 0
1469 + and (max_num_tokens < 0 or num_tokens <= max_num_tokens)
1470 + )
1471
1472 @property
1473 def use_deepep_ht_kernels(self):
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+10 -15
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import (
13 )
14 from vllm.model_executor.layers.fused_moe.moe_output import (
15 UnfinalizedMoEOutput,
16 + convert_flashinfer_moe_output,
17 )
18 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
19 TopKWeightAndReduceNoOP,
@@ -278,13 +279,13 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
279 num_tokens = hidden_states.shape[0]
280 # The runner divides by the token count on the host, so an idle rank's
281 # dummy 0-token forward has to keep the finalized (empty) form.
281 - defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
282 + defer = self.moe_config.should_defer_moe_finalize(num_tokens)
283
284 routing_replay_out = self._maybe_make_routing_replay_buffer(
285 num_tokens=num_tokens,
286 device=hidden_states.device,
287 )
287 - out = flashinfer.fused_moe.trtllm_bf16_moe(
288 + flashinfer_output = flashinfer.fused_moe.trtllm_bf16_moe(
289 routing_logits=router_logits,
290 routing_bias=e_score_correction_bias,
291 hidden_states=hidden_states,
@@ -304,17 +305,11 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
305 routing_replay_out=routing_replay_out,
306 do_finalize=not defer,
307 )
308 + routed_output = convert_flashinfer_moe_output(
309 + flashinfer_output,
310 + do_finalize=not defer,
311 + num_tokens=num_tokens,
312 + top_k=self.topk,
313 + )
314 self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
308 - if defer:
309 - # flashinfer returns a flat permute map; the protocol wants
310 - # [num_tokens, top_k] so consumers can read top_k from its shape.
311 - return UnfinalizedMoEOutput(
312 - gemm2_permuted=out[0],
313 - expert_weights=out[1],
314 - expanded_idx_to_permuted_idx=out[2]
315 - .to(torch.int32)
316 - .view(num_tokens, self.topk),
317 - )
318 - # do_finalize=True yields the finalized states (a bare tensor on some
319 - # FlashInfer versions, a single-element list on others).
320 - return out[0] if isinstance(out, (list, tuple)) else out
315 + return routed_output
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+27 -14
@@ -11,6 +11,10 @@ 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 + convert_flashinfer_moe_output,
17 +)
18 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
19 TopKWeightAndReduceNoOP,
20 )
@@ -215,7 +219,7 @@ class TrtLlmMxfp4ExpertsMonolithic(
219 e_score_correction_bias: torch.Tensor | None = None,
220 routed_scaling_factor: float | None = None,
221 topk_group: int | None = None,
218 - ) -> torch.Tensor:
222 + ) -> torch.Tensor | UnfinalizedMoEOutput:
223 from flashinfer import trtllm_fp4_block_scale_moe
224
225 if a1q_scale is not None:
@@ -225,18 +229,22 @@ class TrtLlmMxfp4ExpertsMonolithic(
229 assert hidden_states.dtype == torch.bfloat16
230 x_quant = hidden_states
231 x_scale = None
228 - output = torch.empty(
229 - *hidden_states.shape[:-1],
230 - self.hidden_dim_unpadded,
231 - dtype=torch.bfloat16,
232 - device=hidden_states.device,
233 - )
232 + num_tokens = hidden_states.shape[0]
233 + defer = self.moe_config.should_defer_moe_finalize(num_tokens)
234 + finalized_output = None
235 + if not defer:
236 + finalized_output = torch.empty(
237 + *hidden_states.shape[:-1],
238 + self.hidden_dim_unpadded,
239 + dtype=torch.bfloat16,
240 + device=hidden_states.device,
241 + )
242
243 routing_replay_out = self._maybe_make_routing_replay_buffer(
236 - num_tokens=hidden_states.shape[0],
244 + num_tokens=num_tokens,
245 device=hidden_states.device,
246 )
239 - trtllm_fp4_block_scale_moe(
247 + flashinfer_output = trtllm_fp4_block_scale_moe(
248 routing_logits=router_logits,
249 routing_bias=e_score_correction_bias,
250 hidden_states=x_quant,
@@ -262,16 +270,21 @@ class TrtLlmMxfp4ExpertsMonolithic(
270 local_num_experts=self.local_num_experts,
271 routed_scaling_factor=routed_scaling_factor,
272 routing_method_type=self.routing_method_type,
265 - do_finalize=True,
273 + do_finalize=not defer,
274 activation_type=self._flashinfer_activation_type(activation),
275 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
268 - output=output,
276 + output=finalized_output,
277 routing_replay_out=routing_replay_out,
278 )
271 - self._maybe_dispatch_routing_replay(
272 - routing_replay_out, num_tokens=hidden_states.shape[0]
279 + routed_output = convert_flashinfer_moe_output(
280 + flashinfer_output,
281 + do_finalize=not defer,
282 + num_tokens=num_tokens,
283 + top_k=self.topk,
284 + finalized_output=finalized_output,
285 )
274 - return output
286 + self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
287 + return routed_output
288
289
290 class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular):
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+10 -13
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
15 )
16 from vllm.model_executor.layers.fused_moe.moe_output import (
17 UnfinalizedMoEOutput,
18 + convert_flashinfer_moe_output,
19 )
20 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
21 TopKWeightAndReduceNoOP,
@@ -551,7 +552,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
552 num_tokens = hidden_states.shape[0]
553 # The runner divides by the token count on the host, so an idle rank's
554 # dummy 0-token forward has to keep the finalized (empty) form.
554 - defer = self.moe_config.use_deferred_moe_finalize and num_tokens > 0
555 + defer = self.moe_config.should_defer_moe_finalize(num_tokens)
556
557 routing_replay_out = self._maybe_make_routing_replay_buffer(
558 num_tokens=num_tokens,
@@ -560,7 +561,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
561 # Invoke kernel.
562 # NOTE: Activation padding and output
563 # truncation are handled by the MoE runner's
563 - result = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
564 + flashinfer_output = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
565 routing_logits=router_logits,
566 routing_bias=e_score_correction_bias,
567 hidden_states=hidden_states,
@@ -594,15 +595,11 @@ class TrtLlmNvFp4ExpertsMonolithic(
595 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
596 routing_replay_out=routing_replay_out,
597 )
598 + routed_output = convert_flashinfer_moe_output(
599 + flashinfer_output,
600 + do_finalize=not defer,
601 + num_tokens=num_tokens,
602 + top_k=self.topk,
603 + )
604 self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens)
598 - if defer:
599 - # flashinfer returns a flat permute map; the protocol wants
600 - # [num_tokens, top_k] so consumers can read top_k from its shape.
601 - return UnfinalizedMoEOutput(
602 - gemm2_permuted=result[0],
603 - expert_weights=result[1],
604 - expanded_idx_to_permuted_idx=result[2]
605 - .to(torch.int32)
606 - .view(num_tokens, self.topk),
607 - )
608 - return result[0]
605 + return routed_output
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+4 -3
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import (
13 FusedMoEParallelConfig,
14 FusedMoEQuantConfig,
15 )
16 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
17 from vllm.model_executor.layers.quantization.base_config import (
18 QuantizeMethodBase,
19 )
@@ -161,7 +162,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
162 shared_experts_input: Input for shared experts (if any)
163
164 Returns:
164 - Output tensor from routed experts
165 + Output tensor from routed experts.
166 """
167 raise NotImplementedError
168
@@ -171,7 +172,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
172 x: torch.Tensor,
173 router_logits: torch.Tensor,
174 input_ids: torch.Tensor | None = None,
174 - ) -> torch.Tensor:
175 + ) -> torch.Tensor | UnfinalizedMoEOutput:
176 """
177 Apply the MoE operation using monolithic kernels.
178
@@ -181,6 +182,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
182 router_logits: Router logits (routing done internally)
183
184 Returns:
184 - Output tensor from routed experts
185 + Finalized routed states or a deferred-finalize output.
186 """
187 raise NotImplementedError
vllm/model_executor/layers/fused_moe/modular_kernel.py
+21 -6
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
22 FusedMoEQuantConfig,
23 RoutingMethodType,
24 )
25 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
26 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
27 SharedExperts,
28 SharedExpertsOrder,
@@ -425,6 +426,9 @@ class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):
426 described above for the monolithic case.
427 """
428
429 + def supports_deferred_moe_finalize(self) -> bool:
430 + return False
431 +
432 @abstractmethod
433 def prepare(
434 self,
@@ -1078,7 +1082,7 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts):
1082 e_score_correction_bias: torch.Tensor | None = None,
1083 routed_scaling_factor: float | None = None,
1084 topk_group: int | None = None,
1081 - ) -> torch.Tensor:
1085 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1086 """
1087 Same as apply(), except uses router_logits as opposed
1088 to the topk_ids and topk_weights. This is useful for kernels
@@ -1548,7 +1552,7 @@ class FusedMoEKernelMonolithicImpl:
1552 e_score_correction_bias: torch.Tensor | None = None,
1553 routed_scaling_factor: float | None = None,
1554 topk_group: int | None = None,
1551 - ) -> torch.Tensor:
1555 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1556 """
1557 Same as forward(), except uses router_logits as opposed
1558 to the topk_ids and topk_weights. This is used for kernels
@@ -1579,9 +1583,14 @@ class FusedMoEKernelMonolithicImpl:
1583 topk_group=topk_group,
1584 )
1585
1582 - output = self.prepare_finalize.finalize(fused_out)
1583 -
1584 - return output
1586 + if isinstance(fused_out, UnfinalizedMoEOutput):
1587 + if not self.prepare_finalize.supports_deferred_moe_finalize():
1588 + raise RuntimeError(
1589 + f"{type(self.prepare_finalize).__name__} cannot pass through "
1590 + "a deferred MoE output."
1591 + )
1592 + return fused_out
1593 + return self.prepare_finalize.finalize(fused_out)
1594
1595
1596 @final
@@ -1664,6 +1673,12 @@ class FusedMoEKernel:
1673 """
1674 return self.prepare_finalize.output_is_reduced()
1675
1676 + def supports_deferred_moe_finalize(self) -> bool:
1677 + return (
1678 + isinstance(self.prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic)
1679 + and self.prepare_finalize.supports_deferred_moe_finalize()
1680 + )
1681 +
1682 def apply_monolithic(
1683 self,
1684 hidden_states: torch.Tensor,
@@ -1679,7 +1694,7 @@ class FusedMoEKernel:
1694 e_score_correction_bias: torch.Tensor | None = None,
1695 routed_scaling_factor: float | None = None,
1696 topk_group: int | None = None,
1682 - ) -> torch.Tensor:
1697 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1698 assert isinstance(self.impl, FusedMoEKernelMonolithicImpl)
1699 return self.impl.apply(
1700 hidden_states=hidden_states,
vllm/model_executor/layers/fused_moe/moe_output.py
+90
@@ -33,6 +33,96 @@ class UnfinalizedMoEOutput:
33 expanded_idx_to_permuted_idx: torch.Tensor
34
35
36 +def convert_flashinfer_moe_output(
37 + flashinfer_output: torch.Tensor | list[torch.Tensor],
38 + *,
39 + do_finalize: bool,
40 + num_tokens: int,
41 + top_k: int,
42 + finalized_output: torch.Tensor | None = None,
43 +) -> torch.Tensor | UnfinalizedMoEOutput:
44 + """Normalize the two FlashInfer TRTLLM MoE return layouts.
45 +
46 + Args:
47 + flashinfer_output: Tensor returned by the FlashInfer BF16 wrapper's
48 + legacy finalized path, or its mode-dependent tensor list.
49 + do_finalize: Whether FlashInfer ran its top-k finalize step.
50 + num_tokens: Number of input tokens.
51 + top_k: Number of routed experts per token.
52 + finalized_output: Optional destination passed to FlashInfer's ``output``
53 + argument.
54 +
55 + Returns:
56 + A finalized tensor or the structured deferred-finalize output.
57 +
58 + Raises:
59 + ValueError: If FlashInfer returns an unexpected layout.
60 + """
61 +
62 + if num_tokens < 0:
63 + raise ValueError("num_tokens must be non-negative.")
64 + if top_k <= 0:
65 + raise ValueError("top_k must be positive.")
66 +
67 + if do_finalize:
68 + if isinstance(flashinfer_output, torch.Tensor):
69 + returned_output = flashinfer_output
70 + elif len(flashinfer_output) == 1 and isinstance(
71 + flashinfer_output[0], torch.Tensor
72 + ):
73 + returned_output = flashinfer_output[0]
74 + else:
75 + raise ValueError(
76 + "Finalized FlashInfer MoE output must contain exactly one tensor."
77 + )
78 + if finalized_output is not None and (
79 + returned_output.shape != finalized_output.shape
80 + or returned_output.dtype != finalized_output.dtype
81 + or returned_output.device != finalized_output.device
82 + or returned_output.data_ptr() != finalized_output.data_ptr()
83 + ):
84 + raise ValueError(
85 + "Finalized FlashInfer MoE output must alias the provided destination."
86 + )
87 + return returned_output
88 +
89 + if finalized_output is not None:
90 + raise ValueError("Deferred FlashInfer MoE output cannot have a destination.")
91 + if isinstance(flashinfer_output, torch.Tensor) or len(flashinfer_output) != 3:
92 + raise ValueError(
93 + "Deferred FlashInfer MoE output must contain GEMM2 output, routing "
94 + "weights, and a permutation map."
95 + )
96 + gemm2_permuted, expert_weights, expanded_idx = flashinfer_output
97 + if gemm2_permuted.ndim != 2:
98 + raise ValueError("Deferred FlashInfer GEMM2 output must be rank-2.")
99 + if expanded_idx.dtype != torch.int32:
100 + raise ValueError("Deferred FlashInfer permutation map must use int32.")
101 + expected_routes = num_tokens * top_k
102 + if expert_weights.numel() != expected_routes:
103 + raise ValueError(
104 + "Deferred FlashInfer routing weights must contain "
105 + f"num_tokens * top_k ({expected_routes}) elements."
106 + )
107 + if expanded_idx.numel() != expected_routes:
108 + raise ValueError(
109 + "Deferred FlashInfer permutation map must contain "
110 + f"num_tokens * top_k ({expected_routes}) elements."
111 + )
112 + if not (gemm2_permuted.device == expert_weights.device == expanded_idx.device):
113 + raise ValueError("Deferred FlashInfer outputs must share a device.")
114 + if not all(
115 + output.is_contiguous()
116 + for output in (gemm2_permuted, expert_weights, expanded_idx)
117 + ):
118 + raise ValueError("Deferred FlashInfer outputs must be contiguous.")
119 + return UnfinalizedMoEOutput(
120 + gemm2_permuted=gemm2_permuted,
121 + expert_weights=expert_weights.view(num_tokens, top_k),
122 + expanded_idx_to_permuted_idx=expanded_idx.view(num_tokens, top_k),
123 + )
124 +
125 +
126 @dataclass
127 class MoEOutput:
128 """A MoE layer's output with its final reduction still open.
vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py
+3
@@ -98,6 +98,9 @@ class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):
98
99
100 class MoEPrepareAndFinalizeNoDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic):
101 + def supports_deferred_moe_finalize(self) -> bool:
102 + return True
103 +
104 @property
105 def activation_format(self) -> mk.FusedMoEActivationFormat:
106 return mk.FusedMoEActivationFormat.Standard
vllm/model_executor/layers/fused_moe/routed_experts.py
+4 -3
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import (
19 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
20 FusedMoEMethodBase,
21 )
22 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
23 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
24 UnquantizedFusedMoEMethod,
25 )
@@ -1216,7 +1217,7 @@ class RoutedExperts(PluggableLayer):
1217 shared_experts_input: Input for shared experts (if any)
1218
1219 Returns:
1219 - Output tensor from routed experts
1220 + Output tensor from routed experts.
1221 """
1222 assert not self.quant_method.is_monolithic
1223
@@ -1235,7 +1236,7 @@ class RoutedExperts(PluggableLayer):
1236 x: torch.Tensor,
1237 router_logits: torch.Tensor | None = None,
1238 input_ids: torch.Tensor | None = None,
1238 - ) -> torch.Tensor:
1239 + ) -> torch.Tensor | UnfinalizedMoEOutput:
1240 """
1241 Execute routed experts using the quantization method's apply function.
1242
@@ -1250,7 +1251,7 @@ class RoutedExperts(PluggableLayer):
1251 input_ids: input ids for DeepSeek V4
1252
1253 Returns:
1253 - Output tensor from routed experts
1254 + Finalized routed states or a deferred-finalize output.
1255 """
1256 assert self.quant_method.is_monolithic
1257
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+42 -20
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Callable, Iterable
4 from contextlib import nullcontext
5 -from typing import TYPE_CHECKING
5 +from typing import TYPE_CHECKING, cast
6
7 import torch
8 import torch.nn.functional as F
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.config import (
28 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
29 FusedMoEMethodBase,
30 )
31 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
32 from vllm.model_executor.layers.fused_moe.routed_experts import (
33 RoutedExperts,
34 )
@@ -121,11 +122,14 @@ def _moe_forward(
122 hidden_dim_unpadded: int,
123 ) -> torch.Tensor:
124 layer = get_layer_from_name(_resolve_layer_name(layer_name))
124 - return layer._forward_impl(
125 - hidden_states,
126 - router_logits,
127 - shared_experts_input,
128 - input_ids,
125 + return cast(
126 + torch.Tensor,
127 + layer._forward_impl(
128 + hidden_states,
129 + router_logits,
130 + shared_experts_input,
131 + input_ids,
132 + ),
133 )
134
135
@@ -155,11 +159,14 @@ def _moe_forward_shared(
159 hidden_dim_unpadded: int,
160 ) -> tuple[torch.Tensor, torch.Tensor]:
161 layer = get_layer_from_name(_resolve_layer_name(layer_name))
158 - return layer._forward_impl(
159 - hidden_states,
160 - router_logits,
161 - shared_experts_input,
162 - input_ids,
162 + return cast(
163 + tuple[torch.Tensor, torch.Tensor],
164 + layer._forward_impl(
165 + hidden_states,
166 + router_logits,
167 + shared_experts_input,
168 + input_ids,
169 + ),
170 )
171
172
@@ -207,8 +214,10 @@ direct_register_custom_op(
214
215
216 def _unpack(
210 - result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
211 -) -> tuple[torch.Tensor | None, torch.Tensor]:
217 + result: torch.Tensor
218 + | UnfinalizedMoEOutput
219 + | tuple[torch.Tensor, torch.Tensor | UnfinalizedMoEOutput],
220 +) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]:
221 if isinstance(result, tuple):
222 return result
223 else:
@@ -573,7 +582,7 @@ class MoERunner(MoERunnerInterface):
582 router_logits: torch.Tensor,
583 shared_experts_input: torch.Tensor | None,
584 input_ids: torch.Tensor | None = None,
576 - ) -> tuple[torch.Tensor | None, torch.Tensor]:
585 + ) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]:
586 """Run expert routing and the fused MoE kernel via the quant method.
587
588 Orchestrates shared expert execution (before/after), expert selection
@@ -733,6 +742,7 @@ class MoERunner(MoERunnerInterface):
742
743 # Extract outputs from result
744 shared_output, fused_output = _unpack(result)
745 + fused_output = cast(torch.Tensor, fused_output)
746
747 if og_hidden_dim_pre_xform is not None:
748 fused_output = fused_output[..., :og_hidden_dim_pre_xform]
@@ -808,18 +818,25 @@ class MoERunner(MoERunnerInterface):
818 def _maybe_combine(
819 self,
820 shared_output: torch.Tensor | None,
811 - hidden_states: torch.Tensor,
812 - ) -> torch.Tensor | tuple[torch.Tensor | None, torch.Tensor]:
821 + hidden_states: torch.Tensor | UnfinalizedMoEOutput,
822 + ) -> (
823 + torch.Tensor
824 + | UnfinalizedMoEOutput
825 + | tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]
826 + ):
827 if self.do_naive_dispatch_combine:
828 hidden_states = get_ep_group().combine(
815 - hidden_states, self.moe_config.is_sequence_parallel
829 + cast(torch.Tensor, hidden_states),
830 + self.moe_config.is_sequence_parallel,
831 )
832
833 if (
834 self.moe_config.pcp_size > 1
835 and not self.moe_config.moe_parallel_config.use_all2all_kernels
836 ):
822 - hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0)
837 + hidden_states = get_pcp_group().reduce_scatter(
838 + cast(torch.Tensor, hidden_states), dim=0
839 + )
840
841 if self.shared_experts is not None:
842 assert shared_output is not None
@@ -833,7 +850,11 @@ class MoERunner(MoERunnerInterface):
850 router_logits: torch.Tensor,
851 shared_experts_input: torch.Tensor | None,
852 input_ids: torch.Tensor | None = None,
836 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
853 + ) -> (
854 + torch.Tensor
855 + | UnfinalizedMoEOutput
856 + | tuple[torch.Tensor, torch.Tensor | UnfinalizedMoEOutput]
857 + ):
858 """Entry point called by the custom op to run the MoE computation.
859
860 Handles pre-dispatch setup (gate application, external shared expert
@@ -844,7 +865,8 @@ class MoERunner(MoERunnerInterface):
865 - fused MoE kernel execution
866 - shared expert computation.
867
847 - Returns a single tensor of combined fused and shared output (if present).
868 + Returns routed output, optionally paired with shared-expert output. A
869 + fused consumer may request the routed output in deferred-finalize form.
870 """
871 # TODO(bnell): this can be removed after MK migration is complete.
872 self.routed_experts._ensure_moe_quant_config_init()
vllm/model_executor/layers/quantization/mxfp4.py
+3 -2
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe import modular_kernel as mk
17 from vllm.model_executor.layers.fused_moe.config import (
18 mxfp4_w4a16_moe_quant_config,
19 )
20 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
21 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
22 TRITON_BACKENDS,
23 Mxfp4MoeBackend,
@@ -451,7 +452,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
452 x: torch.Tensor,
453 router_logits: torch.Tensor,
454 input_ids: torch.Tensor | None = None,
454 - ) -> torch.Tensor:
455 + ) -> torch.Tensor | UnfinalizedMoEOutput:
456 assert self.is_monolithic
457 assert self.moe_kernel is not None
458 return self.moe_kernel.apply_monolithic(
@@ -865,7 +866,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
866 x: torch.Tensor,
867 router_logits: torch.Tensor,
868 input_ids: torch.Tensor | None = None,
868 - ) -> torch.Tensor:
869 + ) -> torch.Tensor | UnfinalizedMoEOutput:
870 assert self.is_monolithic
871 assert self.moe_kernel is not None
872 return self.moe_kernel.apply_monolithic(
vllm/models/kimi_k3/amd/latent_moe_runner.py
+4 -3
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import cast
4 +
5 import torch
6
7 from vllm.distributed import (
@@ -7,7 +9,7 @@ from vllm.distributed import (
9 tensor_model_parallel_all_reduce,
10 )
11 from vllm.logger import init_logger
10 -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack
12 +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
13
14 logger = init_logger(__name__)
15
@@ -139,8 +141,7 @@ class ROCmLatentMoERunner(MoERunner):
141 else 0,
142 )
143
142 - shared_output, fused_output = _unpack(result)
143 - assert shared_output is not None
144 + shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result)
145
146 if og_hidden_dim_pre_xform is not None:
147 fused_output = fused_output[..., :og_hidden_dim_pre_xform]
vllm/models/kimi_k3/nvidia/latent_moe_runner.py
+51 -3
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from enum import IntEnum
4 +from typing import cast
5
6 import torch
7
@@ -11,6 +12,7 @@ from vllm.distributed import (
12 tensor_model_parallel_all_reduce,
13 )
14 from vllm.logger import init_logger
15 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
16 from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner, _unpack
17 from vllm.model_executor.layers.layernorm import RMSNorm
18 from vllm.platforms import current_platform
@@ -103,6 +105,24 @@ class LatentMoERunner(MoERunner):
105 assert transform is not None
106 norm = transform.norm
107 assert norm is not None
108 + from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
109 + TrtLlmMxfp4ExpertsMonolithic,
110 + )
111 + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
112 + TrtLlmNvFp4ExpertsMonolithic,
113 + )
114 +
115 + experts_cls = getattr(self._quant_method, "experts_cls", None)
116 + moe_kernel = self._quant_method.moe_kernel
117 + self.moe_config.defer_moe_finalize = (
118 + (
119 + experts_cls is TrtLlmMxfp4ExpertsMonolithic
120 + or experts_cls is TrtLlmNvFp4ExpertsMonolithic
121 + )
122 + and self.moe_config.hidden_dim == self.moe_config.hidden_dim_unpadded
123 + and moe_kernel is not None
124 + and moe_kernel.supports_deferred_moe_finalize()
125 + )
126 from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import (
127 KimiK3LatentMoETailOp,
128 )
@@ -110,11 +130,24 @@ class LatentMoERunner(MoERunner):
130 op = KimiK3LatentMoETailOp.initialize(
131 hidden_size=transform.up_proj.weight.shape[0],
132 latent_size=norm.weight.shape[0],
133 + experts_per_token=(
134 + self.moe_config.experts_per_token
135 + if self.moe_config.use_deferred_moe_finalize
136 + else 0
137 + ),
138 dtype=norm.weight.dtype,
139 device=norm.weight.device,
140 rms_eps=norm.variance_epsilon,
141 )
142 self._k3_latent_moe_tail_op = op
143 + self.moe_config.defer_moe_finalize_max_num_tokens = (
144 + op.contract.max_num_tokens
145 + if self.moe_config.use_deferred_moe_finalize
146 + else -1
147 + )
148 + else:
149 + self.moe_config.defer_moe_finalize = False
150 + self.moe_config.defer_moe_finalize_max_num_tokens = -1
151
152 def _get_zero_residual(
153 self,
@@ -177,7 +210,7 @@ class LatentMoERunner(MoERunner):
210
211 def _small_batch_tail(
212 self,
180 - fused_output: torch.Tensor,
213 + fused_output: torch.Tensor | UnfinalizedMoEOutput,
214 shared_output: torch.Tensor,
215 trunc_size: int | None,
216 ) -> torch.Tensor:
@@ -304,6 +337,22 @@ class LatentMoERunner(MoERunner):
337 )
338 )
339
340 + if (
341 + og_hidden_dim_pre_xform is None
342 + and self.moe_config.should_defer_moe_finalize(hidden_states.shape[0])
343 + ):
344 + unfinalized = self._forward_impl(
345 + hidden_states,
346 + router_logits,
347 + shared_experts_input,
348 + input_ids,
349 + )
350 + shared_output, fused_output = _unpack(unfinalized)
351 + assert shared_output is not None
352 + assert isinstance(fused_output, UnfinalizedMoEOutput)
353 + result = self._small_batch_tail(fused_output, shared_output, None)
354 + return self._maybe_add_zero_expert_output(result)
355 +
356 result = self._forward_entry(
357 hidden_states,
358 router_logits,
@@ -315,8 +364,7 @@ class LatentMoERunner(MoERunner):
364 else 0,
365 )
366
318 - shared_output, fused_output = _unpack(result)
319 - assert shared_output is not None
367 + shared_output, fused_output = cast(tuple[torch.Tensor, torch.Tensor], result)
368
369 if og_hidden_dim_pre_xform is not None:
370 fused_output = fused_output[..., :og_hidden_dim_pre_xform]
vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py
+160 -23
@@ -16,6 +16,8 @@ import torch.distributed as dist
16 import torch.distributed._symmetric_memory as symm_mem
17 from cutlass import BFloat16, Float32, Int32, Int64, Uint32
18
19 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
20 +
21 from .primitives import (
22 NUM_LAMPORT_BUFFERS,
23 PACKED_BYTES,
@@ -104,6 +106,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
106 fp32_internal: bool = False,
107 include_reduce_scatter: bool = True,
108 include_routed: bool = True,
109 + top_k: int = 0,
110 ):
111 validate_shape(
112 tp_size=tp_size,
@@ -153,6 +156,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
156 self.fp32_internal = fp32_internal
157 self.include_reduce_scatter = include_reduce_scatter
158 self.include_routed = include_routed
159 + self.top_k = top_k
160
161 @cute.jit
162 def __call__(
@@ -171,6 +175,8 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
175 m: Int32,
176 epsilon: Float32,
177 stream: cuda.CUstream,
178 + expert_weights: cute.Tensor,
179 + expanded_idx_to_permuted_idx: cute.Tensor,
180 ):
181 self.kernel(
182 latent_source,
@@ -186,6 +192,8 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
192 shared_peer_ptrs,
193 m,
194 epsilon,
195 + expert_weights,
196 + expanded_idx_to_permuted_idx,
197 ).launch(
198 grid=(self.token_ctas, self.cluster_ctas, self.roles),
199 block=(self.threads, 1, 1),
@@ -211,6 +219,8 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
219 shared_peer_ptrs: cute.Tensor,
220 m: Int32,
221 epsilon: Float32,
222 + expert_weights: cute.Tensor,
223 + expanded_idx_to_permuted_idx: cute.Tensor,
224 ):
225 tidx, _, _ = cute.arch.thread_idx()
226 token_cta, cta_y, role = cute.arch.block_idx()
@@ -244,6 +254,8 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
254 logical_role,
255 cluster_rank,
256 tidx,
257 + expert_weights,
258 + expanded_idx_to_permuted_idx,
259 )
260 token = token + self.token_ctas
261
@@ -269,6 +281,8 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
281 role: Int32,
282 cluster_rank: Int32,
283 tidx: Int32,
284 + expert_weights: cute.Tensor,
285 + expanded_idx_to_permuted_idx: cute.Tensor,
286 ):
287 if role == 0:
288 # ---------------- routed AllReduce + RMSNorm ----------------
@@ -292,15 +306,48 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:
306 )
307 dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2))
308
295 - local_ptr = cute.make_ptr(
296 - BFloat16,
297 - (latent_source.iterator + element_offset).llvm_ptr,
298 - cute.AddressSpace.gmem,
299 - assumed_align=16,
300 - )
301 - local_packed = sanitize_negative_zero(
302 - load_global_u32x4(local_ptr, volatile=False)
303 - )
309 + if cutlass.const_expr(self.top_k > 0):
310 + # finalize topk reduction
311 + local_values = cute.make_rmem_tensor(
312 + cute.make_layout((VEC_BF16,)), BFloat16
313 + )
314 + for element in cutlass.range_constexpr(VEC_BF16):
315 + local_values[element] = BFloat16(0.0)
316 + for slot in cutlass.range_constexpr(self.top_k):
317 + permuted_idx = expanded_idx_to_permuted_idx[token, slot]
318 + if permuted_idx >= Int32(0):
319 + permuted_element = (
320 + Int64(permuted_idx) * self.latent_dim
321 + + Int64(packed_idx) * VEC_BF16
322 + )
323 + permuted_ptr = cute.make_ptr(
324 + BFloat16,
325 + (latent_source.iterator + permuted_element).llvm_ptr,
326 + cute.AddressSpace.gmem,
327 + assumed_align=16,
328 + )
329 + values = packed_u32x4_to_bf16x8(
330 + load_global_u32x4(permuted_ptr, volatile=False)
331 + )
332 + weight = expert_weights[token, slot].to(Float32)
333 + for element in cutlass.range_constexpr(VEC_BF16):
334 + scaled = (values[element].to(Float32) * weight).to(BFloat16)
335 + local_values[element] = (
336 + local_values[element].to(Float32) + scaled.to(Float32)
337 + ).to(BFloat16)
338 + local_packed = sanitize_negative_zero(
339 + bf16x8_to_packed_u32x4(local_values.load())
340 + )
341 + else:
342 + local_ptr = cute.make_ptr(
343 + BFloat16,
344 + (latent_source.iterator + element_offset).llvm_ptr,
345 + cute.AddressSpace.gmem,
346 + assumed_align=16,
347 + )
348 + local_packed = sanitize_negative_zero(
349 + load_global_u32x4(local_ptr, volatile=False)
350 + )
351 multicast_offset = (
352 Int64(current_index) * Int64(bytes_per_buffer)
353 + (
@@ -645,6 +692,8 @@ def _compile_key(
692 fp32_internal: bool,
693 include_reduce_scatter: bool,
694 include_routed: bool,
695 + *,
696 + top_k: int,
697 ):
698 return (
699 torch.accelerator.current_device_index(),
@@ -657,6 +706,7 @@ def _compile_key(
706 fp32_internal,
707 include_reduce_scatter,
708 include_routed,
709 + top_k,
710 )
711
712
@@ -673,6 +723,9 @@ def _runtime_args(
723 shared_flags: torch.Tensor,
724 shared_peer_ptrs: torch.Tensor,
725 rms_eps: float,
726 + *,
727 + expert_weights: torch.Tensor,
728 + expanded_idx_to_permuted_idx: torch.Tensor,
729 ):
730 return (
731 to_cute_dynamic_m(latent_source, mode=0, assumed_align=16),
@@ -686,9 +739,11 @@ def _runtime_args(
739 to_cute(shared_workspace, 16),
740 to_cute(shared_flags, 16),
741 to_cute(shared_peer_ptrs, 16),
689 - Int32(latent_source.shape[0]),
742 + Int32(shared_source.shape[0]),
743 Float32(rms_eps),
744 cuda.CUstream(torch.cuda.current_stream(latent_source.device).cuda_stream),
745 + to_cute_dynamic_m(expert_weights, mode=0, assumed_align=16),
746 + to_cute_dynamic_m(expanded_idx_to_permuted_idx, mode=0, assumed_align=16),
747 )
748
749
@@ -712,6 +767,7 @@ def compile_kernel(
767 fp32_internal: bool,
768 include_reduce_scatter: bool = True,
769 include_routed: bool = True,
770 + top_k: int = 0,
771 ) -> None:
772 """Compile the rank/M specialization without retaining caller tensors."""
773
@@ -725,11 +781,17 @@ def compile_kernel(
781 fp32_internal,
782 include_reduce_scatter,
783 include_routed,
784 + top_k=top_k,
785 )
786 if key in _COMPILED:
787 return
788 device = latent_output.device
732 - latent = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device)
789 + latent_rows = max_m * max(top_k, 1)
790 + latent = torch.empty((latent_rows, latent_dim), dtype=torch.bfloat16, device=device)
791 + expert_weights = torch.empty(
792 + (max_m, max(top_k, 1)), dtype=torch.bfloat16, device=device
793 + )
794 + expanded_idx = torch.empty((max_m, max(top_k, 1)), dtype=torch.int32, device=device)
795 gamma = torch.empty((latent_dim,), dtype=torch.bfloat16, device=device)
796 shared = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device)
797 kernel = AllReduceRMSNormWithReduceScatterEarlyExit(
@@ -742,6 +804,7 @@ def compile_kernel(
804 fp32_internal=fp32_internal,
805 include_reduce_scatter=include_reduce_scatter,
806 include_routed=include_routed,
807 + top_k=top_k,
808 )
809 _COMPILED[key] = cute.compile(
810 kernel,
@@ -758,6 +821,8 @@ def compile_kernel(
821 shared_flags,
822 shared_peer_ptrs,
823 rms_eps,
824 + expert_weights=expert_weights,
825 + expanded_idx_to_permuted_idx=expanded_idx,
826 ),
827 )
828
@@ -785,6 +850,9 @@ def launch(
850 fp32_internal: bool,
851 include_reduce_scatter: bool = True,
852 include_routed: bool = True,
853 + expert_weights: torch.Tensor,
854 + expanded_idx_to_permuted_idx: torch.Tensor,
855 + top_k: int = 0,
856 ) -> None:
857 compile_kernel(
858 rank=rank,
@@ -805,6 +873,7 @@ def launch(
873 fp32_internal=fp32_internal,
874 include_reduce_scatter=include_reduce_scatter,
875 include_routed=include_routed,
876 + top_k=top_k,
877 )
878 _COMPILED[
879 _compile_key(
@@ -817,6 +886,7 @@ def launch(
886 fp32_internal,
887 include_reduce_scatter,
888 include_routed,
889 + top_k=top_k,
890 )
891 ](
892 *_runtime_args(
@@ -832,6 +902,8 @@ def launch(
902 shared_flags,
903 shared_peer_ptrs,
904 rms_eps,
905 + expert_weights=expert_weights,
906 + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
907 )
908 )
909
@@ -851,6 +923,7 @@ class CollectiveKernel:
923 max_token_ctas: int,
924 rms_eps: float,
925 fp32_internal: bool,
926 + top_k: int = 0,
927 ) -> None:
928 validate_shape(
929 tp_size=tp_size,
@@ -866,8 +939,16 @@ class CollectiveKernel:
939 self.max_token_ctas = max_token_ctas
940 self.rms_eps = float(rms_eps)
941 self.fp32_internal = fp32_internal
942 + self.top_k = top_k
943 device = torch.device("cuda", torch.accelerator.current_device_index())
944
945 + self._dummy_expert_weights = torch.empty(
946 + (max_m, max(top_k, 1)), dtype=torch.bfloat16, device=device
947 + )
948 + self._dummy_expanded_idx = torch.empty(
949 + (max_m, max(top_k, 1)), dtype=torch.int32, device=device
950 + )
951 +
952 bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2
953 routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer
954 self._routed_workspace = symm_mem.empty(
@@ -947,38 +1028,91 @@ class CollectiveKernel:
1028 shared_peer_ptrs=self._shared_peer_ptrs,
1029 rms_eps=self.rms_eps,
1030 fp32_internal=fp32_internal,
1031 + top_k=top_k,
1032 )
1033 dist.barrier(group=group, device_ids=[device.index])
1034
1035 def __call__(
1036 self,
955 - latent_source: torch.Tensor,
1037 + latent_source: torch.Tensor | UnfinalizedMoEOutput,
1038 shared_source: torch.Tensor,
1039 gamma: torch.Tensor,
1040 ) -> tuple[torch.Tensor, torch.Tensor]:
959 - if latent_source.ndim != 2 or shared_source.ndim != 2:
960 - raise ValueError("latent_source and shared_source must be rank-2")
961 - m = latent_source.shape[0]
1041 + if shared_source.ndim != 2:
1042 + raise ValueError("shared_source must be rank-2")
1043 + expected: list[tuple[torch.Tensor, tuple[int, ...], str, torch.dtype]]
1044 + if isinstance(latent_source, UnfinalizedMoEOutput):
1045 + if self.top_k <= 0:
1046 + raise ValueError("collective was not configured for top-k finalize")
1047 + gemm2_permuted = latent_source.gemm2_permuted
1048 + expert_weights = latent_source.expert_weights
1049 + expanded_idx = latent_source.expanded_idx_to_permuted_idx
1050 + m = expanded_idx.shape[0]
1051 + expected = [
1052 + (
1053 + gemm2_permuted,
1054 + (gemm2_permuted.shape[0], self.latent_dim),
1055 + "gemm2_permuted",
1056 + torch.bfloat16,
1057 + ),
1058 + (
1059 + expert_weights,
1060 + (m, self.top_k),
1061 + "expert_weights",
1062 + torch.bfloat16,
1063 + ),
1064 + (
1065 + expanded_idx,
1066 + (m, self.top_k),
1067 + "expanded_idx_to_permuted_idx",
1068 + torch.int32,
1069 + ),
1070 + ]
1071 + else:
1072 + if self.top_k > 0:
1073 + raise ValueError("top-k collective requires an unfinalized output")
1074 + if latent_source.ndim != 2:
1075 + raise ValueError("latent_source must be rank-2")
1076 + m = latent_source.shape[0]
1077 + gemm2_permuted = latent_source
1078 + expert_weights = self._dummy_expert_weights
1079 + expanded_idx = self._dummy_expanded_idx
1080 + expected = [
1081 + (
1082 + latent_source,
1083 + (m, self.latent_dim),
1084 + "latent_source",
1085 + torch.bfloat16,
1086 + ),
1087 + ]
1088 device = self._routed_workspace.device
963 - expected = (
964 - (latent_source, (m, self.latent_dim), "latent_source"),
965 - (shared_source, (m, self.hidden_dim), "shared_source"),
966 - (gamma, (self.latent_dim,), "gamma"),
1089 + expected.extend(
1090 + [
1091 + (
1092 + shared_source,
1093 + (m, self.hidden_dim),
1094 + "shared_source",
1095 + torch.bfloat16,
1096 + ),
1097 + (gamma, (self.latent_dim,), "gamma", torch.bfloat16),
1098 + ]
1099 )
968 - for tensor, shape, name in expected:
1100 + for tensor, shape, name, dtype in expected:
1101 if (
1102 tensor.shape != shape
971 - or tensor.dtype != torch.bfloat16
1103 + or tensor.dtype != dtype
1104 or tensor.device != device
1105 or not tensor.is_contiguous()
1106 ):
975 - raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}")
1107 + raise ValueError(
1108 + f"{name} must be contiguous CUDA {dtype} {list(shape)}"
1109 + )
1110 if not 1 <= m <= self.max_m:
1111 raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")
1112
1113 with torch.accelerator.device_index(device.index):
1114 launch(
981 - latent_source,
1115 + gemm2_permuted,
1116 gamma,
1117 self._latent_output,
1118 self._routed_workspace,
@@ -997,6 +1131,9 @@ class CollectiveKernel:
1131 max_m=self.max_m,
1132 max_token_ctas=self.max_token_ctas,
1133 fp32_internal=self.fp32_internal,
1134 + top_k=self.top_k,
1135 + expert_weights=expert_weights,
1136 + expanded_idx_to_permuted_idx=expanded_idx,
1137 )
1138 return (
1139 self._latent_output[:m],
vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py
+29 -14
@@ -9,17 +9,18 @@ import torch
9 import torch.distributed as dist
10
11 from vllm.distributed import get_tp_group
12 +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
13 from vllm.model_executor.warmup.cutedsl_warmup import (
14 CuTeDSLCompileUnit,
15 register_cutedsl_warmup_provider,
16 )
17
17 -_MAX_NUM_TOKENS = 16
18 +_MAX_NUM_TOKENS = 128
19 _SKINNY_MAX_NUM_TOKENS = 5
20 _MMA_TILER_MN = (64, 32)
21 _GEMM_CLUSTER_MN = (1, 8)
22 _B_PRIME_STAGES = 2
22 -_COLLECTIVE_TOKEN_CTAS = 8
23 +_COLLECTIVE_TOKEN_CTAS = 32
24 _LAMPORT_COPY_CTAS = 32
25 _LAMPORT_COPY_THREADS = 224
26 _SUPPORTED_TP_SIZES = (8, 16)
@@ -35,6 +36,7 @@ class KimiK3LatentMoETailContract:
36 latent_size: int
37 max_num_tokens: int
38 rms_eps: float
39 + experts_per_token: int
40
41
42 class KimiK3LatentMoETailOp:
@@ -53,6 +55,7 @@ class KimiK3LatentMoETailOp:
55 dtype: torch.dtype,
56 device: torch.device,
57 rms_eps: float,
58 + experts_per_token: int = 0,
59 ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]:
60 tp = get_tp_group()
61 group = tp.device_group
@@ -72,6 +75,7 @@ class KimiK3LatentMoETailOp:
75 latent_size=latent_size,
76 max_num_tokens=_MAX_NUM_TOKENS,
77 rms_eps=float(rms_eps),
78 + experts_per_token=experts_per_token,
79 ),
80 group,
81 )
@@ -85,10 +89,12 @@ class KimiK3LatentMoETailOp:
89 dtype: torch.dtype,
90 device: torch.device,
91 rms_eps: float,
92 + experts_per_token: int = 0,
93 ) -> "KimiK3LatentMoETailOp":
94 contract, group = cls._contract_and_group(
95 hidden_size=hidden_size,
96 latent_size=latent_size,
97 + experts_per_token=experts_per_token,
98 dtype=dtype,
99 device=device,
100 rms_eps=rms_eps,
@@ -140,6 +146,7 @@ class KimiK3LatentMoETailOp:
146 max_token_ctas=_COLLECTIVE_TOKEN_CTAS,
147 rms_eps=contract.rms_eps,
148 fp32_internal=False,
149 + top_k=contract.experts_per_token,
150 )
151 self._up_projection = AdaptiveUpProjectionKernel(
152 group=group,
@@ -191,7 +198,7 @@ class KimiK3LatentMoETailOp:
198
199 def __call__(
200 self,
194 - routed_output: torch.Tensor,
201 + routed_output: torch.Tensor | UnfinalizedMoEOutput,
202 shared_output: torch.Tensor,
203 rms_weight: torch.Tensor,
204 up_weight: torch.Tensor,
@@ -202,7 +209,12 @@ class KimiK3LatentMoETailOp:
209 rms_weight,
210 up_weight,
211 )
205 - self._up_projection.ensure_compiled(routed_output.shape[0])
212 + num_tokens = (
213 + routed_output.expanded_idx_to_permuted_idx.shape[0]
214 + if isinstance(routed_output, UnfinalizedMoEOutput)
215 + else routed_output.shape[0]
216 + )
217 + self._up_projection.ensure_compiled(num_tokens)
218 latent, shared_shard = self._collective(
219 routed_output,
220 shared_output,
@@ -221,24 +233,27 @@ class KimiK3LatentMoETailOp:
233 )
234 return self._lamport_copy(
235 mailbox,
224 - m=routed_output.shape[0],
236 + m=num_tokens,
237 ).squeeze(0)
238
239 def _validate_inputs(
240 self,
229 - routed_output: torch.Tensor,
241 + routed_output: torch.Tensor | UnfinalizedMoEOutput,
242 shared_output: torch.Tensor,
243 rms_weight: torch.Tensor,
244 up_weight: torch.Tensor,
245 ) -> None:
246 contract = self.contract
235 - if routed_output.ndim != 2:
236 - raise ValueError("routed_output must be a 2D tensor.")
237 - num_tokens = routed_output.shape[0]
238 - if routed_output.shape != (num_tokens, contract.latent_size):
239 - raise ValueError(
240 - f"routed_output must have shape [M, {contract.latent_size}]."
241 - )
247 + if isinstance(routed_output, UnfinalizedMoEOutput):
248 + num_tokens = routed_output.expanded_idx_to_permuted_idx.shape[0]
249 + else:
250 + if routed_output.ndim != 2:
251 + raise ValueError("routed_output must be a 2D tensor.")
252 + num_tokens = routed_output.shape[0]
253 + if routed_output.shape != (num_tokens, contract.latent_size):
254 + raise ValueError(
255 + f"routed_output must have shape [M, {contract.latent_size}]."
256 + )
257 if shared_output.shape != (num_tokens, contract.hidden_size):
258 raise ValueError(
259 f"shared_output must have shape [M, {contract.hidden_size}]."
@@ -256,7 +271,7 @@ class KimiK3LatentMoETailOp:
271 f"{contract.max_num_tokens} tokens."
272 )
273
259 - tensors = (routed_output, shared_output, rms_weight, up_weight)
274 + tensors = (shared_output, rms_weight, up_weight)
275 if any(tensor.device != contract.device for tensor in tensors):
276 raise ValueError("All inputs must be on the contract device.")
277 if any(tensor.dtype != contract.dtype for tensor in tensors):