Support deferred MoE finalize for MXFP4 and FP8 block-scale

[b8zhong]

Motivation

The FlashInfer TRT-LLM MoE kernels can return the unweighted expert outputs instead of performing the top-k weighted combine inside the kernel. When the combine is deferred, it is merged with the shared-expert addition into a single kernel, which removes one kernel launch per MoE layer per decode step.

This mode was previously enabled only for the NVFP4 runner. This PR enables it for the MXFP4 runner and for the FP8 block-scale runner. The FP8 per-tensor runner is not changed.

Modifications

1. The deferred combine is enabled for MXFP4 and FP8 block-scale in addition to NVFP4. Each supported quantization and runner combination accepts the deferred form for exactly one routing format, so the required format is now derived per configuration rather than assumed.

2. The routed scaling factor is handled explicitly. On the MXFP4 path the routing decision is computed in advance, and the MoE kernel therefore does not execute the stage that would apply this factor to the expert weights. The factor is instead passed to the fused combine kernel, which multiplies it into the expert weight already held in a register. No additional kernel is launched. All other callers pass a value of 1.0 and their results are unchanged.

Accuracy

AIME 2026, 30 problems, 8 repeats, 240 samples per configuration.

Configuration	pass@1
Baseline	96.67%
This PR	97.08%

Servers:

# baseline
SGLANG_ENABLE_MOE_DEFERRED_FINALIZE=0 sglang serve --trust-remote-code \
  --model-path deepseek-ai/DeepSeek-V4-Flash-0731 --tp 4 \
  --moe-runner-backend flashinfer_mxfp4 --swa-full-tokens-ratio 0.1 --port 30001

# this PR
sglang serve --trust-remote-code \
  --model-path deepseek-ai/DeepSeek-V4-Flash-0731 --tp 4 \
  --moe-runner-backend flashinfer_mxfp4 --swa-full-tokens-ratio 0.1 --port 30000

Evaluation:

sgl-eval run aime26 --base-url http://127.0.0.1:30001/v1 \
  --n-repeats 8 --temperature 1.0 --top-p 0.95 --thinking \
  --num-threads 240 --max-tokens 131072

sgl-eval run aime26 --base-url http://127.0.0.1:30000/v1 \
  --n-repeats 8 --temperature 1.0 --top-p 0.95 --thinking \
  --num-threads 240 --max-tokens 131072

Speed

DeepSeek-V4-Flash-0731, MXFP4 MoE runner, tensor parallel size 4, B300, input length 1024, output length 1024. One measurement per point.

Batch size	Output tok/s, baseline	Output tok/s, this PR	Change
1	193.09	197.46	+2.3%
4	622.83	632.04	+1.5%
8	1163.31	1177.10	+1.2%
16	2107.59	2105.02	-0.1%
64	6265.74	6272.06	+0.1%
256	15597.53	15246.44	-2.3%

Kernel counts from a decode profile of 20 steps at batch size 1, one tensor parallel rank, 43 MoE layers:

Kernel	Baseline	This PR
In-kernel combine	3.241 ms, 860 launches	not launched
Elementwise add for the shared expert	1.324 ms, 860 launches	not launched
Fused combine and shared-expert add	not launched	3.318 ms, 860 launches

At batch size 256 the same comparison gives 9.717 ms for the two baseline kernels and 6.604 ms for the single fused kernel, a reduction of 3.11 ms per 20 decode steps.

---

CI States

Latest PR Test (Base): ❌ Run #31550766019
Latest PR Test (Extra): ❌ Run #31550765937

[b8zhong]

/rerun-test test_glm52_fp8.py

[b8zhong]

/rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

[github-actions]

Results for /rerun-test test_glm52_fp8.py:

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/8-gpu-models/test_glm52_fp8.py

🚀 8-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/8-gpu-models/test_glm52_fp8.py

[github-actions]

Results for /rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

[b8zhong]

Let's just use #34134. We should rip out the finalize without AR fusion later