Repository navigation
[Feat] FlashInfer: Add NVFP4 W4A16 and SiTU support to MegaMoE - #44036
Merged
Fridge003 merged 6 commits intoOct 16, 2026
Merged
Conversation
zianglih
marked this pull request as ready for review
October 13, 2026 22:02
zianglih
requested review from
Alisehen,
AniZpZ,
BBuf,
Edwardf0t1,
FlamingoPg,
Fridge003,
HaiShaw,
JustinTong0323,
OrangeRedeng,
Ying1123,
b8zhong,
ch-wan,
ispobock,
merrymercy,
mmangkad,
sogalin,
wisclmy0611 and
zijiexia
as code owners
October 13, 2026 22:02
Contributor
Author
|
/tag-and-rerun-ci |
ziang-and
force-pushed
the
feat/flashinfer-megamoe-w4a16-followup
branch
2 times, most recently
from
October 14, 2026 05:45
571dab4 to
e81c6bd
Compare
ziang-and
force-pushed
the
feat/flashinfer-megamoe-w4a16-followup
branch
from
October 14, 2026 19:03
e81c6bd to
7337b0c
Compare
Contributor
Author
|
/rerun-failed-ci |
1 similar comment
Contributor
Author
|
/rerun-failed-ci |
Fridge003
approved these changes
Oct 16, 2026
4 of 5 tasks
This was referenced Oct 16, 2026
ByronHsu
added a commit
to ByronHsu/sglang
that referenced
this pull request
Oct 16, 2026
Resolve the SiTU backend allow-list in modelopt_quant.py: keep upstream's FlashInfer MegaMoE (sgl-project#44036) and this PR's CuTe DSL v2 standard path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
@HumansAnd
Add opt-in NVFP4 W4A16 to FlashInfer MegaMoE for serialized ModelOpt checkpoints
and
nvfp4_online, reusingSGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1. The defaultNVFP4 MegaMoE path remains W4A4.
This is the remaining integration from #39210
on top of merged #43789.
Startup autotuning, cache ownership, native workspace pooling and DP token-count
synchronization are already upstream and are not changed here. The original PR's
branch and history are preserved.
Dependency: the existing FlashInfer 0.7.1 pin includes
FlashInfer #5019.
No Dockerfile or dependency changes.
Modifications
native W4A16 config/preprocessing. Keep packed NVFP4 weights/block scales,
per-expert FP32 weight scales, BF16 activations/EP transport/combine and
post-FC2 routing. Neutralize serialized checkpoint activation scales and
bypass activation quantization for online W4A16. Retain the existing source
precision and excluded-layer behavior for online weight conversion.
W4A4/W4A16. Kimi K3 uses standard TopK and local EP/shared-expert handling;
its GEMM1 alpha/clamp fields map to SiTU beta/linear-beta, separately from
SwiGLU clamps. MXFP8/DeepGEMM reject unwired SiTU. Nemotron ReLU-squared is
unsupported. Requirements remain SM100/SM103 and attention-DP=TP=EP;
W4A16 additionally requires BF16 model dtype and combine.
weight/block-scale allocations, then copy newly prepared data into those
allocations. Share this loader/preparation lifecycle between W4A4 and W4A16. Keep the live
MegaMoE layer, tuned profiles and workspace owner so CUDA graph pointers remain
stable; per-expert runtime scales are staged on each forward. FlashInfer 0.7.1
captures W4A4 input normalization as a scalar: reloading a different input
scale is rejected and requires a server restart. Weight, weight-scale and
per-expert FC1-output scale updates remain supported. The registered
W4A16 serialized reload test covers both cache-flush modes; the existing
storage regression covers both precisions.
weight_scale_2tensors to one,preserving ordinary random-weight initialization. Retain the simple
neutral-scale regression test.
SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0uses external FP32 accumulation. Setting it to
1permits nondeterministicBF16 atomic candidates; tuning can still select external reduction.
and A2A behavior. An unset draft A2A backend now resolves to
nonefor aFlashInfer MegaMoE target, and the NVSHMEM guard accounts for that default.
the registered online NVFP4 GSM8K and serialized weight-reload classes.
No full Kimi K3 E2E or dedicated NVSHMEM environment tests are added.
Single-node NVSHMEM defaults
For single-node FlashInfer MegaMoE, with no speculative decoding or a default/
explicit
nonedraft A2A backend, set these defaults before initialization:Explicit environment settings are preserved. These defaults apply to all
supported FlashInfer MegaMoE formats and select direct peer memory access,
avoiding the uninitialized proxy global-exit state fixed by
NVIDIA/nvshmem commit
af9e2391a849.The closed NVIDIA/nvshmem#110 has
additional background. Remote IB/UCX transport and proxy-based global-exit/
timeout polling are disabled; bootstrap networking is unaffected. All EP GPUs
must support peer memory access. Multi-node jobs and explicit other draft A2A
backends receive no automatic defaults.
Example server command
Four Blackwell GPUs, FlashInfer 0.7.1, online expert-weight conversion:
CUDA_VISIBLE_DEVICES=0,1,2,3 \ SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1 \ SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0 \ SGLANG_FP4_IGNORED_LAYERS='' \ SGLANG_ENABLE_JIT_DEEPGEMM=False \ python -m sglang.launch_server \ --model-path Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \ --quantization nvfp4_online --dtype bfloat16 \ --moe-runner-backend flashinfer_megamoe \ --moe-a2a-backend flashinfer_megamoe \ --tp-size 4 --attn-dp-size 4 --ep-size 4 \ --enable-dp-lm-head \ --cuda-graph-max-bs-decode 64 \ --cuda-graph-backend-prefill disabled \ --max-running-requests 64 \ --chunked-prefill-size 16384 \ --mem-fraction-static 0.7Convert every routed expert for this example; excluded FP8 expert layers cannot
use the NVFP4 MegaMoE runner. With attention DP4, the 16384-token configured
prefill chunk resolves to 4096 tokens per rank. The existing startup autotuner
therefore prepares a 4096-token per-rank prefill profile.
Accuracy Tests
Current source:
7337b0c6aa709132170b4128eb83b8db8c4fb3be, based onupstream main
12dc62310288926b0f38e33ff2d466dee2bf6bb9. This includes theCI-required ancestor
ff8d964cf5b6c463966bca5016018d2fa0836326.All six feature commits are patch-identical after rebase; the scope remains 15 files.
Current-head checks: pre-commit and diff checks passed. The focused CPU
adapter harness passed 29 tests; autotune/cache checks passed 15 tests with one
CUDA-only skip; the neutral-scale regression and accessor source check each
passed. These harnesses use real CPU Torch tensors with serving/GPU imports
stubbed or exact-source AST loading; they do not establish CUDA correctness.
The title also passes the exact upstream title-validation script.
C2 evidence scope: the GPU results below predate this rebase. Units, native,
dummy and online cases ran at
7b173e5261ac63ab5a1d6617a1d2ed626ab4aec1;serialized reload cases ran at
562fe4e70f48de798c9260c71af821b8bad10b5f.Those two tested revisions have identical production Python trees and retained
tests; only documentation and removal of the extra W4A4 test class differ.
The integration patches are unchanged by the rebase, but surrounding upstream
runtime code changed. GPU validation has not been rerun on the current head.
Environment: eight B300 SXM6 AC GPUs, 275040 MiB each, driver 590.48.01;
four GPUs per model/native case. Image
lmsysorg/sglang:nightly-dev-cu13-20261013-7224fce1, amd64 digestsha256:009eafc447f47fc595ea5f9820a9d7a07a9d144f3c8b90ab8ef568f4124475a7.Python 3.12.15, PyTorch 2.14.1+cu130 (CUDA 13.0), CuTe DSL 4.8.0,
NVSHMEM 3.4.5 and nvshmem4py 0.4.0. The image already contained all nine
released FlashInfer 0.7.1 / 0.7.1+cu130 packages. No package install/change or
FlashInfer source build. All 373 Python distributions and six SGLang native
libraries stayed unchanged. Only the source-matched SGLang Rust TreeCore
binding was built into an isolated cache; tests load it with build mode
never.Fixtures: Qwen3 FP8 online checkpoint revision
5a5a776300a41aaa681dd7ff0106608ef2bc90db; NVIDIA Qwen3 NVFP4 serialized checkpointrevision
2538ded2a4edb247b4d2b4a8ba24e44bd4c017c3. Cached snapshots were checkedagainst config/index hashes and four shard sizes each; full weight payloads
were not hashed.
MSE recipe with fast math and E4M3-256, all routed experts converted, memory
fraction 0.7, no MTP. The unchanged registered test uses decode graph max 128,
disabled prefill graphs, 200 examples, 128 evaluation threads and a 512-token
output cap. Effective per-worker prefill is 4096, running cap 3839, KV capacity
1,965,696 and context limit 262144. The earlier example command is a separate
configuration, not the measured accuracy setup.
DP counts
[0,1,1,2],[0,1,2,7],[8,8,8,8],[0,4,9,16]. Two owners/startup contexts; five cold AUTO calls per rank, zero warm calls; pooled refs
2→1→0. Borrowed outputs are cloned immediately. Eager/captured outputs are
finite and bit-exact against a separate fixed-default native session before
and after a visible FC1-scale update. This is not an independent BF16 oracle.
context 2048, per-rank prefill 64, running cap 4, memory fraction 0.4, decode
graphs 1/2/4 and native profiles 1/2/4/64; no MTP. This checks finite output,
not model accuracy.
The measured model configurations resolve as follows (all use TP4/attention-DP4/EP4 and no MTP):
KV/graph release-resume, exact text/token IDs and logprobs within 1e-4. The
output comparison routes to DP0; it does not prove equality across all DP
workers. W4A4 is a local control reusing the registered W4A16 class with its
opt-in set to 0; no additional upstream class.
or speedup claim. Server cases leave the three NVSHMEM variables unset to
exercise production defaults; synthetic native cases set them explicitly.
Raw test results
Warnings were the environment's unknown
asyncio_modepytest setting andupstream import/API deprecations. No C2 unit cases skipped. These GPU results are from the explicit C2 run
on the source revisions listed above. Its final checks confirmed the clean
GPU-tested source, unchanged package/native-library hashes and Rust cache, and
all eight GPUs idle. GitHub CI evaluates the rebased head separately.
Native result records below have
passed=true,bit_exact=true,borrowed_outputs_cloned=true,fc1_scale_mutation_observed=true, capacities[1,2,4,8,16], startup contexts2, and pool refs[2,1,0]in every row.The accuracy harness emits latency and throughput fields above; they are raw
incidental output, not a controlled performance comparison. Disjoint GPU groups
were used concurrently for correctness only. No performance conclusion is drawn.
Initial W4A16 startup on HTTP 31010 collided with the concurrently running
W4A4 server on HTTP 31011: the former's detokenizer and latter's tokenizer
both selected port 31245. No model assertions ran in that attempt. Full logs
were preserved; the retry uses HTTP 34010 with unchanged source, helpers,
checkpoint, settings and assertions. The commands below run sequentially.
Completed reload output:
Public reproduction commands and exact registered settings
These commands reproduce the registered unit selection, online W4A16 accuracy class, and serialized reload checks with the public checkpoints. They do not run a serving performance comparison. Execute the model commands one at a time on four idle Blackwell GPUs; change
CUDA_VISIBLE_DEVICESif a different four-GPU group is assigned.Image and exact source
Use
lmsysorg/sglang:nightly-dev-cu13-20261013-7224fce1(OCI index digestsha256:20d36eeb32df0418a165641b51399fd4dc5239dad57f9b8d1f9974c4c9263ca6). The validation node has eight NVIDIA B300 SXM6 AC GPUs, driver 590.48.01. Each model case uses four GPUs.The image already includes released FlashInfer Python/cubin 0.7.1, the aggregate and six architecture JIT-cache packages 0.7.1+cu130, Torch 2.14.1+cu130, Python 3.12.15, CuTe DSL 4.8.0, NVSHMEM 3.4.5, and nvshmem4py 0.4.0. No package installation or FlashInfer source build is needed. Runtime CuTe DSL specialization can compile during first use.
Run inside that image:
These commands pin the GPU-tested reload revision
562fe4e70f48de798c9260c71af821b8bad10b5fto reproduce the reported C2 evidence. The unit/native/dummy/online runs used7b173e5261ac63ab5a1d6617a1d2ed626ab4aec1, with identical production code and retained tests. For a new run on the rebased PR, setTASK_SOURCE_SHA=7337b0c6aa709132170b4128eb83b8db8c4fb3be; no GPU result is claimed for that revision here.Registered unit selection
Pinned registered-model runner
Save this local runner outside the checkout. It downloads or reuses the exact public snapshot, changes only its model path and port, and invokes the actual registered test. For the W4A4 control it changes the same reload class's environment flag and display label in this process; it adds no upstream test class and keeps the inherited assertions intact.
Online W4A16 GSM8K
The unchanged class uses TP4/attention-DP4/EP4, BF16 model dtype, DP LM head, MegaMoE runner/A2A,
nvfp4_online, decode graph maximum 128, disabled prefill graphs, and static memory fraction 0.7. It evaluates 200 GSM8K examples through the completion API with 128 threads, a 512-token limit, andscore > 0.90.The class sets these actual recipe/precision variables; the runner does not replace them:
The example server configuration in the PR body is separate: its graph/running caps are 64 and configured chunk is 16384. Those values are not substituted into this registered accuracy test.
Serialized W4A16 reload and local W4A4 control
Run sequentially:
Both use serialized NVIDIA NVFP4 weights, BF16 model dtype, TP4/attention-DP4/EP4,
--fp4-gemm-backend flashinfer_cutedsl, MegaMoE runner/A2A,--enable-memory-saver, and disabled prefill graphs.SGLANG_MEMORY_SAVER_CUDA_GRAPH=1is set by the inherited launcher. In-kernel FC2 reduction and per-token activation quantization are both disabled. The W4A16 opt-in is 1 for the registered test and 0 for the local W4A4 control.Both retain
flush_cache=TrueandFalsewithabort_all_requests=False, weight/KV/graph release-resume, exact text/token-ID comparisons, and log-probability tolerance1e-4. Generation uses"The capital of France is", temperature 0, and at most 16 new tokens, explicitly routed to DP rank 0. These are same-checkpoint reload tests; they do not establish changed-checkpoint accuracy or all-rank output equality. A W4A4 checkpoint changing the captured input activation scale requires an engine restart.No full Kimi K3/GLM-5.2 E2E, online RL reload, or speedup claim is included.
Speed Tests and Profiling
No controlled serving benchmark or speed comparison was run for this PR.
No speedup is claimed.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #38620506291
Latest PR Test (Extra): ❌ Run #38620505630
Latest PR Test (AMD ROCm 10): ❌ Run #38620505417