[NVIDIA] Support flashinfer Mega Moe - #31470
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
github-actions
Bot
added
documentation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
This depends on FlashInfer PR that is not merged yet. Should we mark this as Draft for now? |
github-actions
Bot
added
the
dependencies
2fbffd3 to
ce4f3ae
Compare
|
/rerun-failed-ci bypass-fastfail 30a27fa |
|
|
||
| The first mega layer's first forward creates it (collective; safe because | ||
| warmup runs the same layer on all ranks in lockstep); later layers reuse it. | ||
| """ |
There was a problem hiding this comment.
Do we have any suggestions on what a fixed flashinfer API would look like?
# Conflicts: # python/sglang/srt/layers/quantization/fp8.py
| ) | ||
|
|
||
|
|
||
| def _ensure_shared_workspace(mega: Any) -> None: |
There was a problem hiding this comment.
[suggestion] _ensure_shared_workspace and warmup_all_flashinfer_megamoe_layers carry multi-paragraph history (a prior commit hash, why FlashInfer's pool missed, the capture repro narrative). That belongs in the PR body; the line only needs the live constraint (share by geometry because per-layer alpha tensors are identity-keyed; build every MegaMOE layer before capture).
Suggestion: Collapse each to one or two sentences stating the cross-module constraint, and drop the changelog/repro story.
| The first mega layer's first forward creates it (collective; safe because | ||
| warmup runs the same layer on all ranks in lockstep); later layers reuse it. | ||
| """ | ||
| if getattr(mega, "_workspace", None) is not None: |
There was a problem hiding this comment.
please do not use getattr/hasattr
|
All NV pipelines passed. |
# Conflicts: # python/sglang/srt/runtime_context.py
Co-authored-by: djns99 <40156487+djns99@users.noreply.github.com> Co-authored-by: 云挚 <ningyunxiao.nyx@antgroup.com> Co-authored-by: Yangmin Li <yangminl@nvidia.com> Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
7 participants
Fork from https://github.com/djns99/sglang/tree/djns99/mega_moe_flashinfer
@djns99 is the main author of this PR.
Motivation
Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciMoE backend benchmark:
flashinfer_trtllm_routed(baseline) vsflashinfer_megamoeBenchmark of the DeepSeek-V4-Flash NVFP4 model comparing two MoE backend configurations
on the serving path.
flashinfer_trtllm_routed(trtllm routed MoE runner + flashinferA2A) is used as the baseline;
flashinfer_megamoeis the compared configuration.All deltas are reported relative to the trtllm_routed baseline.
Environment
nvidia/DeepSeek-V4-Flash-NVFP4attn_dp_size=4, so per-rank decode batch =max_running_requests / 4)Server commands
Baseline —
flashinfer_trtllm_routed:Compared —
flashinfer_megamoe:Client command
Results
Output token throughput (tok/s, higher is better)
Mean TPOT (ms, lower is better)
Full metrics
Takeaways
flashinfer_trtllm_routedwins on throughput — +23% at c32 and +11% at c128 output token throughput, with comparable or better median TPOT. It also sustains higher achieved concurrency (28.8/116.1 vs 26.5/103.7), i.e. it drains the request queue faster.flashinfer_trtllm_routedis the better default for latency-sensitive / low-to-mid concurrency serving; at saturation the choice is throughput-neutral.FlashInfer MegaMOE env knobs
This PR exposes
SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPEfor the FlashInfer NVFP4 MegaMOE path. It selects the cross-rank combine wire format passed toNvfp4CutedslMegaMoeConfig.combine_dtype.Supported values:
bf16: default, exact combine path.mxfp8: quantized combine traffic with a small accuracy tradeoff.nvfp4: lower combine traffic with a small accuracy tradeoff.Example:
mxfp8andnvfp4are incompatible withSGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=1; server args validation now rejects that combination early.MoE backend comparison — DeepSeek-V4-Flash-NVFP4, TP4/DP4/EP4
Baseline:
trtllm. Positive throughput and negative latency are both wins.Tokens/rank = per-DP-rank prefill chunk (DP attention routes each request wholly to one rank;
--chunked-prefill-sizeis divided bydp_sizeinternally).Prefill
2048 tokens/rank (input len 2048, concurrency 128)
Change vs
trtllm:8192 tokens/rank (input len 8192, 640 requests, 2,665,479 input tokens, concurrency 128)
Change vs
trtllm:Decode
Input len 32, output len 1024, 10,240 requests, concurrency 8192, 5,256,337 generated tokens.
Change vs
trtllm:Takeaways
(+18.6% throughput, -16.6% mean TTFT) but only edges trtllm at 8192 tok/rank (+2.1%).
flashinfer_cutedsl_megamoe is the reverse: +9% at 2048, +11.9% at 8192.
plain < ikr < combine_nvfp4. At 2048 tok/rank,plainandcombine_nvfp4are tied(+8.85% vs +9.06%, mean TTFT within 2 ms) and
ikris weakest (+6.19%).(8,620-8,886 tok/s), far tighter than the prefill spread. At concurrency 8192 this workload is
memory-bound, so the MoE kernel is not the bottleneck.
by 4.1%, at the cost of 1.9% throughput — the only backend that trades throughput for tail
latency. flashinfer improves both modestly.
(+1.17% throughput, -2.34% mean TPOT), reproducing the 8192 tok/rank prefill ordering.
flashinfer_cutedsl_megamoe_combine_nvfp4for large-prefill anddecode-heavy serving;
deepgemm megamoeif small-prefill throughput or decode P99 dominates.Reproduce
<VARIANT-SELECTOR>below is whatever selects plain / ikr / combine_nvfp4 — it is not aserver arg (all three record
moe_runner_backend=flashinfer_megamoe).Decode — flashinfer_cutedsl_megamoe
The trtllm baseline is the same command with
--moe-runner-backend flashinfer_trtllm_routed --moe-a2a-backend flashinfer.Decode
Input len 32, output len 1024, 10,240 requests, concurrency 8192, 5,256,337 generated tokens.
Change vs
trtllm:Takeaways
(+18.6% throughput, -16.6% mean TTFT) but only edges trtllm at 8192 tok/rank (+2.1%).
flashinfer_cutedsl_megamoe is the reverse: +9% at 2048, +11.9% at 8192.
plain < ikr < combine_nvfp4. At 2048 tok/rank,plainandcombine_nvfp4are tied(+8.85% vs +9.06%, mean TTFT within 2 ms) and
ikris weakest (+6.19%).(8,620-8,886 tok/s), far tighter than the prefill spread. At concurrency 8192 this workload is
memory-bound, so the MoE kernel is not the bottleneck.
by 4.1%, at the cost of 1.9% throughput — the only backend that trades throughput for tail
latency. flashinfer improves both modestly.
(+1.17% throughput, -2.34% mean TPOT), reproducing the 8192 tok/rank prefill ordering.
flashinfer_cutedsl_megamoe_combine_nvfp4for large-prefill anddecode-heavy serving;
deepgemm megamoeif small-prefill throughput or decode P99 dominates.Reproduce
<VARIANT-SELECTOR>below is whatever selects plain / ikr / combine_nvfp4 — it is not aserver arg (all three record
moe_runner_backend=flashinfer_megamoe).Decode — flashinfer_cutedsl_megamoe
The trtllm baseline is the same command with
--moe-runner-backend flashinfer_trtllm_routed --moe-a2a-backend flashinfer.Prefill 8192 tokens/rank — flashinfer_cutedsl_megamoe
For the 2048 tokens/rank point, use
--chunked-prefill-size 8192 --max-prefill-tokens 8192on theserver and
--random-input-len 2048 --num-prompts 1280on the client.Setup & caveats
kv_cache_dtype=fp8_e4m3,page_size=256, EP4, dp-attention enabledBest run per metric across repeats (max throughput, min latency)
Repeat counts vary: prefill flashinfer variants 2-6 runs each, trtllm/deepgemm 1-2 runs;
decode is a single run per backend
Prefill 2048 tok/rank: the plain/combine_nvfp4 throughput gap (~200 tok/s) is inside the
run-to-run spread (~1,100 tok/s), so treat those two as tied rather than ranked
Decode: single run per backend. The flashinfer variants span only 0.6% in throughput and 0.8%
in mean TPOT, so their relative ordering is suggestive, not established. The deepgemm P99 TPOT
win (-13.4%) is large enough to be credible from one run.
Server commands above are reconstructed from the recorded
server_args;mem_fraction_staticwas auto-derived (0.8 decode, 0.681 / 0.876 prefill) and is not passed explicitly
FLASHINFER_MOE_EP_KNOB_CACHE point to the autotuned configs.
Remaining soft spot: every decode number is a single run, and the three FlashInfer variants differ by well under 1% in throughput. If you want the decode ranking to
hold up in review, 2–3 repeats per variant would settle it; the deepgemm P99 result is the only decode finding that stands on its own.
Accuracy — GSM8K (8-shot, 1316 questions)
No accuracy regression. All three flashinfer_cutedsl_megamoe variants match or slightly exceed
the trtllm baseline, with zero invalid outputs. At n=1316 the binomial standard error is ~0.6 pt,
so the +0.3 to +0.8 pt differences are within noise — the variants should be read as accuracy-
equivalent to the baseline, not better than it.
Update 7/29/2026 — latest 2048 tokens/rank prefill rerun
Workload: DeepSeek-V4-Flash-NVFP4, 2048 tokens/rank prefill, output len 1,
640 requests, max concurrency 128.
Takeaway: in this latest setup, the autotuned
combine_nvfp4config improves throughput and averageTTFT over the non-autotuned
combine_nvfp4path (+13.8% input throughput, -14.2% mean TTFT).Against
deepgemm megamoe, autotunedcombine_nvfp4has slightly higher input throughput (+2.5%)and lower mean TTFT (-4.1%), while
deepgemm megamoehas better tail latency (-10.2% P95 TTFT,-11.8% P99 TTFT).
CI States
CI States
Latest PR Test (Base): ✅ Run #34431504439
Latest PR Test (Extra): ❌ Run #34431504350
Latest PR Test (AMD ROCm 10): ❌ Run #34431504462