@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 97a668152b

[RL Infra][FlashInfer] Enable router replay output from FlashInfer monolithic MoE kernel (#44214)

Signed-off-by: Xuanyu Zhang <xuanyu.zhang@mistral.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com>

Xuanyu Zhang committed Jul 21, 2026 at 01:45 UTC 97a668152b9033caf602b37b88289f03e8faffd8
11 files changed +1099 -12
tests/kernels/moe/test_routed_experts_capture_monolithic.py new
+880
@@ -0,0 +1,880 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""End-to-end tests for routed-expert capture on the monolithic MoE path.
4 +
5 +These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
6 +expert IDs picked by FlashInfer's fused router-and-experts kernels (the
7 +"monolithic" path). When ``set_capture_fn`` is installed on
8 +a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
9 +should:
10 +
11 + * allocate an int16 ``(num_tokens, top_k)`` buffer,
12 + * pass it to FlashInfer as ``routing_replay_out``,
13 + * invoke the callback after the kernel returns.
14 +"""
15 +
16 +from __future__ import annotations
17 +
18 +from types import SimpleNamespace
19 +from unittest.mock import patch
20 +
21 +import pytest
22 +import torch
23 +
24 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
25 +from vllm.model_executor.layers.fused_moe.config import (
26 + FusedMoEConfig,
27 + FusedMoEParallelConfig,
28 + FusedMoEQuantConfig,
29 + RoutingMethodType,
30 + fp8_w8a8_moe_quant_config,
31 +)
32 +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
33 + TrtLlmBf16ExpertsMonolithic,
34 +)
35 +from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
36 + TrtLlmFp8ExpertsMonolithic,
37 +)
38 +from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
39 + TrtLlmNvFp4ExpertsMonolithic,
40 +)
41 +from vllm.platforms import current_platform
42 +
43 +try:
44 + from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
45 +except ImportError:
46 + pytest.skip("flashinfer not available", allow_module_level=True)
47 +
48 +if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
49 + pytest.skip(
50 + "Requires FlashInfer TRT-LLM fused MoE on CUDA",
51 + allow_module_level=True,
52 + )
53 +
54 +if not current_platform.has_device_capability(100):
55 + pytest.skip(
56 + "TRT-LLM fused MoE kernels require SM100+",
57 + allow_module_level=True,
58 + )
59 +
60 +
61 +def _shuffle_bf16_weights_block_major_k(
62 + w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
63 +) -> torch.Tensor:
64 + """Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
65 + ``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
66 + row shuffle.
67 + """
68 + from flashinfer import shuffle_matrix_a
69 + from flashinfer.fused_moe import convert_to_block_layout
70 +
71 + num_experts = w.shape[0]
72 + shuffled = []
73 + for i in range(num_experts):
74 + t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
75 + shuffled.append(convert_to_block_layout(t, block_k))
76 + return torch.stack(shuffled).view(torch.bfloat16)
77 +
78 +
79 +def _make_bf16_monolithic_experts(
80 + num_experts: int,
81 + top_k: int,
82 + hidden_size: int,
83 + intermediate_size: int,
84 + routing_method: RoutingMethodType,
85 + device: torch.device,
86 +) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
87 + """Construct the monolithic BF16 experts plus the BlockMajorK weights
88 + expected by ``trtllm_bf16_moe``.
89 + """
90 + parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
91 + moe_config = FusedMoEConfig(
92 + num_experts=num_experts,
93 + experts_per_token=top_k,
94 + hidden_dim=hidden_size,
95 + intermediate_size=intermediate_size,
96 + num_local_experts=num_experts,
97 + num_logical_experts=num_experts,
98 + moe_parallel_config=parallel_cfg,
99 + in_dtype=torch.bfloat16,
100 + activation=MoEActivation.SILU,
101 + device=device,
102 + routing_method=routing_method,
103 + max_num_tokens=max(8, 1),
104 + )
105 + quant_config = FusedMoEQuantConfig.make(
106 + quant_dtype=None,
107 + per_act_token_quant=False,
108 + per_out_ch_quant=False,
109 + block_shape=None,
110 + )
111 +
112 + experts = TrtLlmBf16ExpertsMonolithic(
113 + moe_config=moe_config, quant_config=quant_config
114 + )
115 +
116 + gemm1 = (
117 + torch.randn(
118 + num_experts,
119 + 2 * intermediate_size,
120 + hidden_size,
121 + device=device,
122 + dtype=torch.bfloat16,
123 + )
124 + * 0.1
125 + )
126 + gemm2 = (
127 + torch.randn(
128 + num_experts,
129 + hidden_size,
130 + intermediate_size,
131 + device=device,
132 + dtype=torch.bfloat16,
133 + )
134 + * 0.1
135 + )
136 + w13 = _shuffle_bf16_weights_block_major_k(gemm1)
137 + w2 = _shuffle_bf16_weights_block_major_k(gemm2)
138 + return experts, w13, w2
139 +
140 +
141 +def _run_bf16_monolithic(
142 + experts: TrtLlmBf16ExpertsMonolithic,
143 + hidden_states: torch.Tensor,
144 + w13: torch.Tensor,
145 + w2: torch.Tensor,
146 + router_logits: torch.Tensor,
147 + num_experts: int,
148 + *,
149 + n_group: int | None = None,
150 + topk_group: int | None = None,
151 + routed_scaling_factor: float | None = None,
152 + e_score_correction_bias: torch.Tensor | None = None,
153 +) -> torch.Tensor:
154 + return experts.apply(
155 + hidden_states=hidden_states,
156 + w1=w13,
157 + w2=w2,
158 + router_logits=router_logits,
159 + activation=MoEActivation.SILU,
160 + global_num_experts=num_experts,
161 + expert_map=None,
162 + a1q_scale=None,
163 + apply_router_weight_on_input=False,
164 + num_expert_group=n_group,
165 + topk_group=topk_group,
166 + e_score_correction_bias=e_score_correction_bias,
167 + routed_scaling_factor=routed_scaling_factor,
168 + )
169 +
170 +
171 +_DSV3_NUM_EXPERTS = 32
172 +_DSV3_N_GROUP = 4
173 +_DSV3_TOPK_GROUP = 2
174 +
175 +
176 +def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
177 + return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
178 +
179 +
180 +@pytest.mark.parametrize("num_tokens", [2, 7, 16])
181 +@pytest.mark.parametrize("top_k", [2, 4])
182 +def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
183 + num_tokens: int,
184 + top_k: int,
185 +) -> None:
186 + """The capture callback should receive the int16 routed-expert IDs the
187 + kernel actually used, the values should be valid expert indices, and
188 + each token should pick ``top_k`` distinct experts."""
189 + if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
190 + pytest.skip(
191 + f"DSV3 requires top_k <= n_group * topk_group "
192 + f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
193 + )
194 + torch.manual_seed(0)
195 + device = torch.device("cuda:0")
196 +
197 + num_experts = _DSV3_NUM_EXPERTS
198 + hidden_size = 1024
199 + intermediate_size = 1024
200 +
201 + experts, w13, w2 = _make_bf16_monolithic_experts(
202 + num_experts=num_experts,
203 + top_k=top_k,
204 + hidden_size=hidden_size,
205 + intermediate_size=intermediate_size,
206 + routing_method=RoutingMethodType.DeepSeekV3,
207 + device=device,
208 + )
209 +
210 + captured: list[torch.Tensor] = []
211 +
212 + def capture_fn(replay_out: torch.Tensor) -> None:
213 + captured.append(replay_out.clone())
214 +
215 + assert experts.supports_routing_replay_capture()
216 + experts.set_capture_fn(capture_fn)
217 +
218 + hidden_states = (
219 + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
220 + )
221 + router_logits = torch.rand(
222 + num_tokens, num_experts, device=device, dtype=torch.float32
223 + )
224 + routing_bias = _make_dsv3_routing_bias(num_experts, device)
225 +
226 + _ = _run_bf16_monolithic(
227 + experts,
228 + hidden_states=hidden_states,
229 + w13=w13,
230 + w2=w2,
231 + router_logits=router_logits,
232 + num_experts=num_experts,
233 + n_group=_DSV3_N_GROUP,
234 + topk_group=_DSV3_TOPK_GROUP,
235 + routed_scaling_factor=1.0,
236 + e_score_correction_bias=routing_bias,
237 + )
238 +
239 + assert len(captured) == 1
240 + replay = captured[0]
241 + assert replay.dtype == torch.int16
242 + assert replay.shape == (num_tokens, top_k)
243 + assert (replay >= 0).all(), f"got out-of-range values: {replay}"
244 + assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
245 +
246 + for t in range(num_tokens):
247 + unique = replay[t].unique()
248 + assert unique.numel() == top_k, (
249 + f"token {t}: expected {top_k} distinct experts, "
250 + f"got {unique.numel()} ({replay[t].tolist()})"
251 + )
252 +
253 +
254 +@pytest.mark.parametrize("num_tokens", [2, 7, 16])
255 +@pytest.mark.parametrize(
256 + "routing_method",
257 + [
258 + RoutingMethodType.Renormalize,
259 + RoutingMethodType.RenormalizeNaive,
260 + ],
261 +)
262 +def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
263 + num_tokens: int,
264 + routing_method: RoutingMethodType,
265 +) -> None:
266 + """Routing replay works for non-DeepSeekV3 routing methods too.
267 + FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
268 + torch.manual_seed(0)
269 + device = torch.device("cuda:0")
270 +
271 + num_experts = 8
272 + top_k = 2
273 + hidden_size = 1024
274 + intermediate_size = 1024
275 +
276 + experts, w13, w2 = _make_bf16_monolithic_experts(
277 + num_experts=num_experts,
278 + top_k=top_k,
279 + hidden_size=hidden_size,
280 + intermediate_size=intermediate_size,
281 + routing_method=routing_method,
282 + device=device,
283 + )
284 +
285 + captured: list[torch.Tensor] = []
286 + experts.set_capture_fn(lambda r: captured.append(r.clone()))
287 +
288 + hidden_states = (
289 + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
290 + )
291 + router_logits = torch.rand(
292 + num_tokens, num_experts, device=device, dtype=torch.float32
293 + )
294 +
295 + _ = _run_bf16_monolithic(
296 + experts,
297 + hidden_states=hidden_states,
298 + w13=w13,
299 + w2=w2,
300 + router_logits=router_logits,
301 + num_experts=num_experts,
302 + )
303 +
304 + assert len(captured) == 1
305 + replay = captured[0]
306 + assert replay.dtype == torch.int16
307 + assert replay.shape == (num_tokens, top_k)
308 + assert (replay >= 0).all(), f"got out-of-range values: {replay}"
309 + assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
310 + for t in range(num_tokens):
311 + unique = replay[t].unique()
312 + assert unique.numel() == top_k, (
313 + f"token {t}: expected {top_k} distinct experts, "
314 + f"got {unique.numel()} ({replay[t].tolist()})"
315 + )
316 +
317 +
318 +def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
319 + """With no callback installed the kernel should not see a
320 + ``routing_replay_out`` tensor — verify the helper short-circuits."""
321 + torch.manual_seed(0)
322 + device = torch.device("cuda:0")
323 + experts, _, _ = _make_bf16_monolithic_experts(
324 + num_experts=_DSV3_NUM_EXPERTS,
325 + top_k=2,
326 + hidden_size=1024,
327 + intermediate_size=1024,
328 + routing_method=RoutingMethodType.DeepSeekV3,
329 + device=device,
330 + )
331 + # No callback installed.
332 + buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
333 + assert buf is None
334 +
335 + # Dispatch is also a no-op.
336 + experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
337 +
338 +
339 +def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
340 + """FlashInfer's ``routing_replay_out`` is supported by all routing
341 + methods, so ``supports_routing_replay_capture`` should be True
342 + regardless of routing method."""
343 + device = torch.device("cuda:0")
344 + for routing_method in (
345 + RoutingMethodType.DeepSeekV3,
346 + RoutingMethodType.Renormalize,
347 + RoutingMethodType.RenormalizeNaive,
348 + ):
349 + experts, _, _ = _make_bf16_monolithic_experts(
350 + num_experts=_DSV3_NUM_EXPERTS,
351 + top_k=2,
352 + hidden_size=1024,
353 + intermediate_size=1024,
354 + routing_method=routing_method,
355 + device=device,
356 + )
357 + assert experts.supports_routing_replay_capture() is True, (
358 + f"{routing_method!r} should support routing replay capture"
359 + )
360 +
361 +
362 +def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
363 + """When capture is installed, the allocated buffer is int16 and shaped
364 + ``(num_tokens, experts_per_token)``."""
365 + device = torch.device("cuda:0")
366 + experts, _, _ = _make_bf16_monolithic_experts(
367 + num_experts=_DSV3_NUM_EXPERTS,
368 + top_k=4,
369 + hidden_size=1024,
370 + intermediate_size=1024,
371 + routing_method=RoutingMethodType.DeepSeekV3,
372 + device=device,
373 + )
374 + experts.set_capture_fn(lambda r: None)
375 + buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
376 + assert buf is not None
377 + assert buf.dtype == torch.int16
378 + assert buf.shape[0] >= 11
379 + assert buf.shape[1] == 4
380 + assert buf.device.type == "cuda"
381 +
382 +
383 +def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
384 + """End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
385 + on the monolithic experts and verify the captured rows land in the
386 + capturer's device buffer at the correct layer slot.
387 +
388 + Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
389 + for the monolithic path: a single closure is installed on the monolithic
390 + ``fused_experts`` (in addition to ``router.set_capture_fn`` on the
391 + non-monolithic path) and the capturer routes per-layer based on the
392 + closed-over ``layer_id``.
393 + """
394 + from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
395 + RoutedExpertsCapturer,
396 + )
397 +
398 + torch.manual_seed(7)
399 + device = torch.device("cuda:0")
400 + num_tokens = 4
401 + top_k = 2
402 + num_experts = _DSV3_NUM_EXPERTS
403 + hidden_size = 1024
404 + intermediate_size = 1024
405 +
406 + experts, w13, w2 = _make_bf16_monolithic_experts(
407 + num_experts=num_experts,
408 + top_k=top_k,
409 + hidden_size=hidden_size,
410 + intermediate_size=intermediate_size,
411 + routing_method=RoutingMethodType.DeepSeekV3,
412 + device=device,
413 + )
414 +
415 + num_layers = 3
416 + layer_id = 1
417 + capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
418 + capturer.dp_rank = 0
419 + capturer.tp_size = 1
420 + capturer.device_buffer = torch.full(
421 + (num_tokens + 4, num_layers, top_k),
422 + -1,
423 + dtype=torch.int32,
424 + device=device,
425 + )
426 +
427 + def capture_fn(replay_out: torch.Tensor) -> None:
428 + capturer.capture(layer_id, replay_out)
429 +
430 + experts.set_capture_fn(capture_fn)
431 +
432 + hidden_states = (
433 + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
434 + )
435 + router_logits = torch.rand(
436 + num_tokens, num_experts, device=device, dtype=torch.float32
437 + )
438 + routing_bias = _make_dsv3_routing_bias(num_experts, device)
439 +
440 + # Patch get_forward_context to return a dp_metadata=None context so the
441 + # capturer takes the single-DP branch.
442 + import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
443 +
444 + with patch.object(
445 + rec,
446 + "get_forward_context",
447 + return_value=SimpleNamespace(dp_metadata=None),
448 + ):
449 + _ = _run_bf16_monolithic(
450 + experts,
451 + hidden_states=hidden_states,
452 + w13=w13,
453 + w2=w2,
454 + router_logits=router_logits,
455 + num_experts=num_experts,
456 + n_group=_DSV3_N_GROUP,
457 + topk_group=_DSV3_TOPK_GROUP,
458 + routed_scaling_factor=1.0,
459 + e_score_correction_bias=routing_bias,
460 + )
461 +
462 + captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
463 + # Valid expert IDs at this layer.
464 + assert (captured >= 0).all()
465 + assert (captured < num_experts).all()
466 + for t in range(num_tokens):
467 + unique = captured[t].unique()
468 + assert unique.numel() == top_k, (
469 + f"token {t}: expected {top_k} distinct experts at layer "
470 + f"{layer_id}, got {unique.numel()}"
471 + )
472 +
473 + # Other layers / trailing token rows untouched.
474 + for other_layer in range(num_layers):
475 + if other_layer == layer_id:
476 + continue
477 + assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
478 + f"layer {other_layer} should be untouched, got writes"
479 + )
480 + assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
481 + "tail rows beyond num_tokens should remain sentinel"
482 + )
483 +
484 +
485 +# ----------------------------------------------------------------------------
486 +# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
487 +# ----------------------------------------------------------------------------
488 +
489 +
490 +def _make_fp8_block_scale_monolithic_experts(
491 + num_experts: int,
492 + top_k: int,
493 + hidden_size: int,
494 + intermediate_size: int,
495 + device: torch.device,
496 +) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
497 + """Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
498 + code path with DSV3 routing.
499 +
500 + Weights are shuffled into the BlockMajorK layout the kernel expects
501 + (same helper the vLLM weight loader uses for DeepSeek-FP8 models).
502 + """
503 + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
504 + _shuffle_deepseek_fp8_moe_weights,
505 + )
506 +
507 + block_k = 128
508 + parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
509 + moe_config = FusedMoEConfig(
510 + num_experts=num_experts,
511 + experts_per_token=top_k,
512 + hidden_dim=hidden_size,
513 + intermediate_size=intermediate_size,
514 + num_local_experts=num_experts,
515 + num_logical_experts=num_experts,
516 + moe_parallel_config=parallel_cfg,
517 + in_dtype=torch.bfloat16,
518 + activation=MoEActivation.SILU,
519 + device=device,
520 + routing_method=RoutingMethodType.DeepSeekV3,
521 + max_num_tokens=max(8, 1),
522 + )
523 +
524 + # Random fp8 weights + ones-block scales (the kernel decoder only cares
525 + # that the per-block scales are present and finite for routing/replay).
526 + gemm1 = torch.randn(
527 + num_experts, 2 * intermediate_size, hidden_size, device=device
528 + ).to(torch.float8_e4m3fn)
529 + gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
530 + torch.float8_e4m3fn
531 + )
532 + w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
533 +
534 + w1_scale = torch.ones(
535 + num_experts,
536 + 2 * intermediate_size // block_k,
537 + hidden_size // block_k,
538 + device=device,
539 + dtype=torch.float32,
540 + )
541 + w2_scale = torch.ones(
542 + num_experts,
543 + hidden_size // block_k,
544 + intermediate_size // block_k,
545 + device=device,
546 + dtype=torch.float32,
547 + )
548 + quant_config = fp8_w8a8_moe_quant_config(
549 + w1_scale=w1_scale,
550 + w2_scale=w2_scale,
551 + block_shape=[block_k, block_k],
552 + per_act_token_quant=False,
553 + )
554 +
555 + experts = TrtLlmFp8ExpertsMonolithic(
556 + moe_config=moe_config, quant_config=quant_config
557 + )
558 + return experts, w13_shuffled, w2_shuffled
559 +
560 +
561 +def _run_fp8_block_scale_monolithic(
562 + experts: TrtLlmFp8ExpertsMonolithic,
563 + hidden_states_fp8: torch.Tensor,
564 + hidden_states_scale: torch.Tensor,
565 + w13: torch.Tensor,
566 + w2: torch.Tensor,
567 + router_logits: torch.Tensor,
568 + num_experts: int,
569 + routing_bias: torch.Tensor,
570 +) -> torch.Tensor:
571 + return experts.apply(
572 + hidden_states=hidden_states_fp8,
573 + w1=w13,
574 + w2=w2,
575 + router_logits=router_logits,
576 + activation=MoEActivation.SILU,
577 + global_num_experts=num_experts,
578 + expert_map=None,
579 + # The block-scale apply path reads ``a1q_scale`` and transposes it
580 + # to ``(hidden_size/128, num_tokens)`` for the kernel call.
581 + a1q_scale=hidden_states_scale,
582 + apply_router_weight_on_input=False,
583 + num_expert_group=_DSV3_N_GROUP,
584 + topk_group=_DSV3_TOPK_GROUP,
585 + e_score_correction_bias=routing_bias,
586 + routed_scaling_factor=1.0,
587 + )
588 +
589 +
590 +@pytest.mark.parametrize("num_tokens", [2, 7, 16])
591 +@pytest.mark.parametrize("top_k", [2, 4])
592 +def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
593 + num_tokens: int,
594 + top_k: int,
595 +) -> None:
596 + """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
597 + path, DSV3 routing) captures valid expert IDs."""
598 + if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
599 + pytest.skip(
600 + f"DSV3 requires top_k <= n_group * topk_group "
601 + f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
602 + )
603 + torch.manual_seed(0)
604 + device = torch.device("cuda:0")
605 +
606 + num_experts = _DSV3_NUM_EXPERTS
607 + hidden_size = 1024
608 + intermediate_size = 1024
609 +
610 + experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
611 + num_experts=num_experts,
612 + top_k=top_k,
613 + hidden_size=hidden_size,
614 + intermediate_size=intermediate_size,
615 + device=device,
616 + )
617 + assert experts.supports_routing_replay_capture()
618 +
619 + captured: list[torch.Tensor] = []
620 + experts.set_capture_fn(lambda r: captured.append(r.clone()))
621 +
622 + # Per-token / per-block hidden scales (ones is fine for the routing
623 + # path; the GEMM output isn't being asserted on).
624 + hidden_states = (
625 + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
626 + ).to(torch.float8_e4m3fn)
627 + hidden_states_scale = torch.ones(
628 + num_tokens, hidden_size // 128, device=device, dtype=torch.float32
629 + )
630 + router_logits = torch.rand(
631 + num_tokens, num_experts, device=device, dtype=torch.float32
632 + )
633 + routing_bias = _make_dsv3_routing_bias(num_experts, device)
634 +
635 + _ = _run_fp8_block_scale_monolithic(
636 + experts,
637 + hidden_states_fp8=hidden_states,
638 + hidden_states_scale=hidden_states_scale,
639 + w13=w13,
640 + w2=w2,
641 + router_logits=router_logits,
642 + num_experts=num_experts,
643 + routing_bias=routing_bias,
644 + )
645 +
646 + assert len(captured) == 1
647 + replay = captured[0]
648 + assert replay.dtype == torch.int16
649 + assert replay.shape == (num_tokens, top_k)
650 + assert (replay >= 0).all(), f"got out-of-range values: {replay}"
651 + assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
652 + for t in range(num_tokens):
653 + unique = replay[t].unique()
654 + assert unique.numel() == top_k, (
655 + f"token {t}: expected {top_k} distinct experts, "
656 + f"got {unique.numel()} ({replay[t].tolist()})"
657 + )
658 +
659 +
660 +# ----------------------------------------------------------------------------
661 +# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
662 +# ----------------------------------------------------------------------------
663 +
664 +
665 +def _make_nvfp4_monolithic_experts(
666 + num_experts: int,
667 + top_k: int,
668 + hidden_size: int,
669 + intermediate_size: int,
670 + device: torch.device,
671 +) -> tuple[
672 + TrtLlmNvFp4ExpertsMonolithic,
673 + torch.Tensor, # w13 (packed nvfp4 uint8)
674 + torch.Tensor, # w13 block-scale (fp8)
675 + torch.Tensor, # w2 (packed nvfp4 uint8)
676 + torch.Tensor, # w2 block-scale (fp8)
677 + torch.Tensor, # input global scale (per-tensor float32)
678 +]:
679 + """Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
680 + and DSV3 routing.
681 +
682 + NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
683 + "global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
684 + ``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
685 + * weights: uint8 (packed fp4) ``(E, M, K//2)``
686 + * weight scales: fp8 ``(E, M, K//16)``
687 + * hidden states: uint8 (packed fp4) ``(N, K//2)``
688 + * hidden state scales: fp8 ``(N, K//16)``
689 + """
690 + from flashinfer import fp4_quantize
691 +
692 + block_size = 16
693 + parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
694 + moe_config = FusedMoEConfig(
695 + num_experts=num_experts,
696 + experts_per_token=top_k,
697 + hidden_dim=hidden_size,
698 + intermediate_size=intermediate_size,
699 + num_local_experts=num_experts,
700 + num_logical_experts=num_experts,
701 + moe_parallel_config=parallel_cfg,
702 + in_dtype=torch.bfloat16,
703 + activation=MoEActivation.SILU,
704 + device=device,
705 + routing_method=RoutingMethodType.DeepSeekV3,
706 + max_num_tokens=max(8, 1),
707 + )
708 +
709 + gemm1 = torch.randn(
710 + num_experts,
711 + 2 * intermediate_size,
712 + hidden_size,
713 + device=device,
714 + dtype=torch.bfloat16,
715 + )
716 + gemm2 = torch.randn(
717 + num_experts,
718 + hidden_size,
719 + intermediate_size,
720 + device=device,
721 + dtype=torch.bfloat16,
722 + )
723 + # Per-tensor weight scaling factor (used to build ``g1_alphas`` /
724 + # ``g2_alphas`` below).
725 + w_global_scale = torch.tensor(1.0, device=device)
726 + # Per-tensor input scaling factor.
727 + a_global_scale = torch.tensor(1.0, device=device)
728 +
729 + w13_q, w13_scale = fp4_quantize(
730 + gemm1,
731 + w_global_scale,
732 + block_size,
733 + sf_use_ue8m0=False,
734 + is_sf_swizzled_layout=False,
735 + )
736 + w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
737 + num_experts, 2 * intermediate_size, hidden_size // block_size
738 + )
739 + w2_q, w2_scale = fp4_quantize(
740 + gemm2,
741 + w_global_scale,
742 + block_size,
743 + sf_use_ue8m0=False,
744 + is_sf_swizzled_layout=False,
745 + )
746 + w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
747 + num_experts, hidden_size, intermediate_size // block_size
748 + )
749 +
750 + # NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
751 + # g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
752 + g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
753 + a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
754 +
755 + quant_config = FusedMoEQuantConfig.make(
756 + quant_dtype="nvfp4",
757 + per_act_token_quant=False,
758 + per_out_ch_quant=False,
759 + block_shape=None,
760 + w1_scale=w13_scale,
761 + w2_scale=w2_scale,
762 + g1_alphas=g_alphas,
763 + g2_alphas=g_alphas,
764 + a1_gscale=a_global_scale,
765 + a2_gscale=a2_gscale,
766 + )
767 +
768 + experts = TrtLlmNvFp4ExpertsMonolithic(
769 + moe_config=moe_config, quant_config=quant_config
770 + )
771 + return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
772 +
773 +
774 +def _run_nvfp4_monolithic(
775 + experts: TrtLlmNvFp4ExpertsMonolithic,
776 + hidden_states_q: torch.Tensor,
777 + hidden_states_scale: torch.Tensor,
778 + router_logits: torch.Tensor,
779 + num_experts: int,
780 + routing_bias: torch.Tensor,
781 +) -> torch.Tensor:
782 + """The monolithic NVFP4 apply expects packed fp4 hidden states + the
783 + matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
784 + # Stash the weight tensors on the experts in the locations the apply()
785 + # implementation reads from (it pulls them from quant_config / scales
786 + # already; w1/w2 come in as args).
787 + return experts.apply(
788 + hidden_states=hidden_states_q,
789 + w1=experts._w13_packed,
790 + w2=experts._w2_packed,
791 + router_logits=router_logits,
792 + activation=MoEActivation.SILU,
793 + global_num_experts=num_experts,
794 + expert_map=None,
795 + a1q_scale=hidden_states_scale,
796 + apply_router_weight_on_input=False,
797 + num_expert_group=_DSV3_N_GROUP,
798 + topk_group=_DSV3_TOPK_GROUP,
799 + e_score_correction_bias=routing_bias,
800 + routed_scaling_factor=1.0,
801 + )
802 +
803 +
804 +@pytest.mark.parametrize("num_tokens", [2, 7, 16])
805 +@pytest.mark.parametrize("top_k", [2, 4])
806 +def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
807 + num_tokens: int,
808 + top_k: int,
809 +) -> None:
810 + """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
811 + on the DSV3 routing path."""
812 + if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
813 + pytest.skip(
814 + f"DSV3 requires top_k <= n_group * topk_group "
815 + f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
816 + )
817 + from flashinfer import fp4_quantize
818 +
819 + torch.manual_seed(0)
820 + device = torch.device("cuda:0")
821 +
822 + num_experts = _DSV3_NUM_EXPERTS
823 + hidden_size = 1024
824 + intermediate_size = 1024
825 + block_size = 16
826 +
827 + experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
828 + num_experts=num_experts,
829 + top_k=top_k,
830 + hidden_size=hidden_size,
831 + intermediate_size=intermediate_size,
832 + device=device,
833 + )
834 + # The apply() reads w1/w2 from its args, but we keep them on the experts
835 + # for convenience of the helper.
836 + experts._w13_packed = w13_q
837 + experts._w2_packed = w2_q
838 +
839 + assert experts.supports_routing_replay_capture()
840 + captured: list[torch.Tensor] = []
841 + experts.set_capture_fn(lambda r: captured.append(r.clone()))
842 +
843 + hidden_states = (
844 + torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
845 + )
846 + hidden_states_q, hidden_states_scale = fp4_quantize(
847 + hidden_states,
848 + a_gs,
849 + block_size,
850 + sf_use_ue8m0=False,
851 + is_sf_swizzled_layout=False,
852 + )
853 + # The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
854 + # ``hidden_states_scale`` in its native (uint8 packed) form.
855 + router_logits = torch.rand(
856 + num_tokens, num_experts, device=device, dtype=torch.float32
857 + )
858 + routing_bias = _make_dsv3_routing_bias(num_experts, device)
859 +
860 + _ = _run_nvfp4_monolithic(
861 + experts,
862 + hidden_states_q=hidden_states_q,
863 + hidden_states_scale=hidden_states_scale,
864 + router_logits=router_logits,
865 + num_experts=num_experts,
866 + routing_bias=routing_bias,
867 + )
868 +
869 + assert len(captured) == 1
870 + replay = captured[0]
871 + assert replay.dtype == torch.int16
872 + assert replay.shape == (num_tokens, top_k)
873 + assert (replay >= 0).all(), f"got out-of-range values: {replay}"
874 + assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
875 + for t in range(num_tokens):
876 + unique = replay[t].unique()
877 + assert unique.numel() == top_k, (
878 + f"token {t}: expected {top_k} distinct experts, "
879 + f"got {unique.numel()} ({replay[t].tolist()})"
880 + )
tests/model_executor/test_routed_experts_capture.py
+55
@@ -66,6 +66,12 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
66 )
67
68
69 +def _make_modular_routed_experts():
70 + return types.SimpleNamespace(
71 + quant_method=types.SimpleNamespace(is_monolithic=False),
72 + )
73 +
74 +
75 def test_base_router_capture_pre_eplb_mapping():
76 router = _make_router()
77 captured = []
@@ -122,6 +128,8 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
128 def __init__(self):
129 self.layer_id = 7
130 self.router = _make_router()
131 + self.routed_experts = _make_modular_routed_experts()
132 + self._quant_method = self.routed_experts.quant_method
133
134 class DummyCapturer:
135 def __init__(self):
@@ -160,6 +168,8 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
168 def __init__(self):
169 self.layer_id = 11
170 self.router = _make_router()
171 + self.routed_experts = _make_modular_routed_experts()
172 + self._quant_method = self.routed_experts.quant_method
173
174 class DummyCapturer:
175 def __init__(self):
@@ -197,6 +207,8 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
207 def __init__(self, layer_id):
208 self.layer_id = layer_id
209 self.router = _make_router()
210 + self.routed_experts = _make_modular_routed_experts()
211 + self._quant_method = self.routed_experts.quant_method
212
213 target_module = DummyFusedMoE(layer_id=7)
214 draft_module = DummyFusedMoE(layer_id=0)
@@ -222,6 +234,49 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
234 assert draft_module.router.capture_fn is None
235
236
237 +def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
238 + from vllm.v1.worker import gpu_model_runner as gmr
239 +
240 + class DummyFusedMoE:
241 + def __init__(self):
242 + self.layer_id = 3
243 + self.router = _make_router()
244 + # Use a concrete monolithic expert and override its capability
245 + # instead of instantiating the abstract base class directly.
246 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
247 + CPUExpertsFp8,
248 + )
249 +
250 + fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
251 + self.routed_experts = types.SimpleNamespace(
252 + quant_method=types.SimpleNamespace(
253 + is_monolithic=True,
254 + moe_kernel=types.SimpleNamespace(
255 + impl=types.SimpleNamespace(fused_experts=fused_experts)
256 + ),
257 + )
258 + )
259 + self._quant_method = self.routed_experts.quant_method
260 + self._quant_method.moe_kernel.impl.fused_experts = fused_experts
261 + fused_experts.supports_routing_replay_capture = lambda: False
262 +
263 + class DummyCapturer:
264 + def capture(self, layer_id, topk_ids):
265 + pass
266 +
267 + dummy_module = DummyFusedMoE()
268 + import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
269 +
270 + monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
271 +
272 + dummy_self = types.SimpleNamespace(
273 + model=types.SimpleNamespace(modules=lambda: [dummy_module])
274 + )
275 +
276 + with pytest.raises(ValueError, match="monolithic MoE kernel"):
277 + gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
278 +
279 +
280 def test_routed_experts_capturer_single_dp_no_metadata():
281 """dp_metadata is None: capture writes the full topk_ids rows."""
282 capturer = _capturer_with_buffer(dp_rank=0)
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+13 -1
@@ -34,6 +34,9 @@ class TrtLlmBf16ExpertsBase:
34 monolithic interfaces.
35 """
36
37 + def supports_routing_replay_capture(self) -> bool:
38 + return True
39 +
40 def __init__(
41 self,
42 moe_config: FusedMoEConfig,
@@ -246,7 +249,11 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
249
250 assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
251
249 - return flashinfer.fused_moe.trtllm_bf16_moe(
252 + routing_replay_out = self._maybe_make_routing_replay_buffer(
253 + num_tokens=hidden_states.shape[0],
254 + device=hidden_states.device,
255 + )
256 + out = flashinfer.fused_moe.trtllm_bf16_moe(
257 routing_logits=router_logits,
258 routing_bias=e_score_correction_bias,
259 hidden_states=hidden_states,
@@ -263,4 +270,9 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
270 routing_method_type=self.routing_method_type,
271 activation_type=activation_to_flashinfer_int(activation),
272 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
273 + routing_replay_out=routing_replay_out,
274 + )
275 + self._maybe_dispatch_routing_replay(
276 + routing_replay_out, num_tokens=hidden_states.shape[0]
277 )
278 + return out
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+23 -1
@@ -265,6 +265,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
265 Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.
266 """
267
268 + def supports_routing_replay_capture(self) -> bool:
269 + return True
270 +
271 def __init__(
272 self,
273 moe_config: FusedMoEConfig,
@@ -403,6 +406,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
406 n_group = num_expert_group or 0
407 selected_topk_group = topk_group or 0
408
409 + routing_replay_out = self._maybe_make_routing_replay_buffer(
410 + num_tokens=hidden_states.shape[0],
411 + device=hidden_states.device,
412 + )
413 +
414 kwargs = dict(
415 routing_logits=router_logits,
416 routing_bias=e_score_correction_bias,
@@ -427,11 +435,16 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
435 use_shuffled_weight=use_shuffled_weight,
436 weight_layout=weight_layout,
437 fp8_quantization_type=fp8_quant_type,
438 + routing_replay_out=routing_replay_out,
439 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
440 )
441 if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
442 kwargs["activation_type"] = activation_type
434 - return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
443 + result = flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
444 + self._maybe_dispatch_routing_replay(
445 + routing_replay_out, num_tokens=hidden_states.shape[0]
446 + )
447 + return result
448
449 def _apply_per_tensor(
450 self,
@@ -464,6 +477,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
477 else:
478 assert not apply_router_weight_on_input
479
480 + routing_replay_out = self._maybe_make_routing_replay_buffer(
481 + num_tokens=hidden_states.shape[0],
482 + device=hidden_states.device,
483 + )
484 +
485 out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe(
486 routing_logits=router_logits,
487 routing_bias=e_score_correction_bias,
@@ -485,6 +503,10 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
503 routing_method_type=self.routing_method_type,
504 activation_type=activation_type,
505 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
506 + routing_replay_out=routing_replay_out,
507 + )
508 + self._maybe_dispatch_routing_replay(
509 + routing_replay_out, num_tokens=hidden_states.shape[0]
510 )
511 return out
512
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+11 -1
@@ -122,6 +122,9 @@ class TrtLlmMxfp4ExpertsMonolithic(
122 Wraps flashinfer.trtllm_fp4_block_scale_moe().
123 """
124
125 + def supports_routing_replay_capture(self) -> bool:
126 + return True
127 +
128 @staticmethod
129 def _supports_parallel_config(
130 moe_parallel_config: FusedMoEParallelConfig,
@@ -184,6 +187,10 @@ class TrtLlmMxfp4ExpertsMonolithic(
187 device=hidden_states.device,
188 )
189
190 + routing_replay_out = self._maybe_make_routing_replay_buffer(
191 + num_tokens=hidden_states.shape[0],
192 + device=hidden_states.device,
193 + )
194 trtllm_fp4_block_scale_moe(
195 routing_logits=router_logits.to(torch.bfloat16),
196 routing_bias=None,
@@ -213,8 +220,11 @@ class TrtLlmMxfp4ExpertsMonolithic(
220 do_finalize=True,
221 tune_max_num_tokens=max(self.max_capture_size, 1),
222 output=output,
223 + routing_replay_out=routing_replay_out,
224 + )
225 + self._maybe_dispatch_routing_replay(
226 + routing_replay_out, num_tokens=hidden_states.shape[0]
227 )
217 -
228 return output
229
230
vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
+14 -1
@@ -122,6 +122,9 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
122 # The kernel handles quantization internally.
123 return True
124
125 + def supports_routing_replay_capture(self) -> bool:
126 + return True
127 +
128 def apply(
129 self,
130 hidden_states: torch.Tensor,
@@ -144,7 +147,12 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
147
148 assert self.w1_scale is not None
149 assert self.w2_scale is not None
147 - return flashinfer_trtllm_mxint4_moe(
150 +
151 + routing_replay_out = self._maybe_make_routing_replay_buffer(
152 + num_tokens=hidden_states.shape[0],
153 + device=hidden_states.device,
154 + )
155 + result = flashinfer_trtllm_mxint4_moe(
156 x=hidden_states,
157 router_logits=router_logits,
158 w13_weight_packed=w1,
@@ -160,4 +168,9 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
168 topk_group=topk_group,
169 e_score_correction_bias=e_score_correction_bias,
170 routing_method_type=self.routing_method,
171 + routing_replay_out=routing_replay_out,
172 + )
173 + self._maybe_dispatch_routing_replay(
174 + routing_replay_out, num_tokens=hidden_states.shape[0]
175 )
176 + return result
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+13 -1
@@ -435,6 +435,9 @@ class TrtLlmNvFp4ExpertsMonolithic(
435 Monolithic version of the kernel (router + experts).
436 """
437
438 + def supports_routing_replay_capture(self) -> bool:
439 + return True
440 +
441 @staticmethod
442 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
443 """The modular implementation should be used for the Dp/Ep or EPLB case."""
@@ -509,10 +512,14 @@ class TrtLlmNvFp4ExpertsMonolithic(
512
513 output1_scale_gate_scalar = self.quant_config.g1_alphas
514
515 + routing_replay_out = self._maybe_make_routing_replay_buffer(
516 + num_tokens=hidden_states.shape[0],
517 + device=hidden_states.device,
518 + )
519 # Invoke kernel.
520 # NOTE: Activation padding and output
521 # truncation are handled by the MoE runner's
515 - return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
522 + result = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
523 routing_logits=router_logits,
524 routing_bias=e_score_correction_bias,
525 hidden_states=hidden_states,
@@ -544,4 +551,9 @@ class TrtLlmNvFp4ExpertsMonolithic(
551 activation_type=activation_to_flashinfer_int(activation),
552 per_token_scale=per_token_scale,
553 tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
554 + routing_replay_out=routing_replay_out,
555 )[0]
556 + self._maybe_dispatch_routing_replay(
557 + routing_replay_out, num_tokens=hidden_states.shape[0]
558 + )
559 + return result
vllm/model_executor/layers/fused_moe/modular_kernel.py
+51
@@ -1000,6 +1000,57 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts):
1000 def is_monolithic() -> bool:
1001 return True
1002
1003 + routing_replay_capture_fn: Callable[[torch.Tensor], None] | None = None
1004 + _routing_replay_buffer: torch.Tensor | None = None
1005 +
1006 + def supports_routing_replay_capture(self) -> bool:
1007 + """Whether this expert supports routing replay capture.
1008 +
1009 + Subclasses backed by a kernel that exposes routed expert IDs
1010 + (e.g. FlashInfer's ``routing_replay_out``) should override.
1011 + """
1012 + return False
1013 +
1014 + def set_capture_fn(
1015 + self,
1016 + capture_fn: Callable[[torch.Tensor], None] | None,
1017 + ) -> None:
1018 + self.routing_replay_capture_fn = capture_fn
1019 + if capture_fn is None:
1020 + self._routing_replay_buffer = None
1021 + return
1022 + self._routing_replay_buffer = torch.empty(
1023 + (self.moe_config.max_num_tokens, self.moe_config.experts_per_token),
1024 + dtype=torch.int16,
1025 + device=self.moe_config.device,
1026 + )
1027 +
1028 + def _maybe_make_routing_replay_buffer(
1029 + self,
1030 + num_tokens: int,
1031 + device: torch.device,
1032 + ) -> torch.Tensor | None:
1033 + if self.routing_replay_capture_fn is None:
1034 + return None
1035 + buf = self._routing_replay_buffer
1036 + assert buf is not None
1037 + if buf.shape[0] < num_tokens or buf.device != device:
1038 + raise ValueError(
1039 + "Routing replay buffer was initialized for "
1040 + f"{buf.shape[0]} tokens on {buf.device}, but the kernel "
1041 + f"received {num_tokens} tokens on {device}."
1042 + )
1043 + return buf
1044 +
1045 + def _maybe_dispatch_routing_replay(
1046 + self,
1047 + routing_replay_out: torch.Tensor | None,
1048 + num_tokens: int,
1049 + ) -> None:
1050 + if routing_replay_out is None or self.routing_replay_capture_fn is None:
1051 + return
1052 + self.routing_replay_capture_fn(routing_replay_out[:num_tokens])
1053 +
1054 def apply(
1055 self,
1056 hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+11 -1
@@ -90,10 +90,20 @@ class RoutedExpertsCapturer:
90 ) -> None:
91 hf_config = vllm_config.model_config.hf_text_config
92 num_experts_per_tok = _get_num_experts_per_tok(hf_config)
93 + num_layers = hf_config.num_hidden_layers
94 + logger.info(
95 + "RoutedExpertsCapturer: allocating buffer with "
96 + "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
97 + "(hf_config.model_type=%s)",
98 + max_num_batched_tokens,
99 + num_layers,
100 + num_experts_per_tok,
101 + getattr(hf_config, "model_type", "unknown"),
102 + )
103 self.device_buffer = torch.zeros(
104 (
105 max_num_batched_tokens,
96 - hf_config.num_hidden_layers,
106 + num_layers,
107 num_experts_per_tok,
108 ),
109 # Use int32 for the device / host transit buffers: it
vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
+2
@@ -190,6 +190,7 @@ def flashinfer_trtllm_mxint4_moe(
190 topk_group: int | None = None,
191 e_score_correction_bias: torch.Tensor | None = None,
192 routing_method_type: int | None = None,
193 + routing_replay_out: torch.Tensor | None = None,
194 ) -> torch.Tensor:
195 """
196 Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
@@ -262,6 +263,7 @@ def flashinfer_trtllm_mxint4_moe(
263 do_finalize=True,
264 output=None,
265 tune_max_num_tokens=8192,
266 + routing_replay_out=routing_replay_out,
267 )
268 if isinstance(out, (tuple, list)):
269 out = out[0]
vllm/v1/worker/gpu_model_runner.py
+26 -6
@@ -7676,17 +7676,37 @@ class GPUModelRunner(
7676
7677 def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:
7678 from vllm.model_executor.layers.fused_moe.layer import MoERunner
7679 + from vllm.model_executor.layers.fused_moe.modular_kernel import (
7680 + FusedMoEExpertsMonolithic,
7681 + )
7682 from vllm.model_executor.layers.fused_moe.router.base_router import (
7683 BaseRouter,
7684 )
7685
7686 for module in self.model.modules():
7684 - if isinstance(module, MoERunner) and isinstance(module.router, BaseRouter):
7685 - layer_id = module.layer_id
7686 -
7687 - def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):
7688 - _capturer.capture(_layer_id, topk_ids)
7689 -
7687 + if not isinstance(module, MoERunner):
7688 + continue
7689 + layer_id = module.layer_id
7690 +
7691 + def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):
7692 + _capturer.capture(_layer_id, topk_ids)
7693 +
7694 + quant_method = module._quant_method
7695 + moe_kernel = getattr(quant_method, "moe_kernel", None)
7696 + impl = getattr(moe_kernel, "impl", None)
7697 + fused_experts = getattr(impl, "fused_experts", None)
7698 + if quant_method.is_monolithic:
7699 + if not (
7700 + isinstance(fused_experts, FusedMoEExpertsMonolithic)
7701 + and fused_experts.supports_routing_replay_capture()
7702 + ):
7703 + raise ValueError(
7704 + "--enable-return-routed-experts is not supported with "
7705 + f"monolithic MoE kernel {type(fused_experts).__name__}; "
7706 + "routed expert IDs would be silently all-zero."
7707 + )
7708 + fused_experts.set_capture_fn(_capture_fn)
7709 + elif isinstance(module.router, BaseRouter):
7710 module.router.set_capture_fn(_capture_fn)
7711
7712 def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: