Repository navigation
Conversation
为单节点 NVFP4 W4A16 设置 P2P-only NVSHMEM 默认值,绕过 proxy 未初始化退出状态;保留显式配置并排除使用其他 A2A 后端的 draft。按要求不增加 behavioral regression tests。
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 NVFP4 W4A16 to the existing FlashInfer MegaMoE backend for serialized ModelOpt
checkpoints and
nvfp4_online. ReuseSGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1from #35120; the default NVFP4
MegaMoE path remains W4A4.
Merge prerequisite: SGLang must first bump FlashInfer to
0.7.1, which
includes FlashInfer #5019.
The current 0.7.0.post1 pin lacks W4A16 and uses the old W4A4 FC2 keyword.
Dockerfiles and dependency pins are outside this PR's scope.
Modifications
and use FlashInfer's native config, preprocessing and pooled workspace APIs.
W4A16 uses packed NVFP4 weights/block scales, per-expert FP32 weight scales,
BF16 expert activations/EP transport/combine and post-FC2 routing. Checkpoint
activation scales are neutralized; online quantization retains its source
precision and excluded-layer semantics. No new opt-in or quantization interface.
synchronized DP token counts and graph buckets while expert inputs remain local.
Kimi K3 NVFP4 W4A4/W4A16 uses standard TopK and local EP/shared-expert handling;
SiTU maps GEMM1 alpha/clamp fields to beta/linear-beta, separately from SwiGLU
clamps. MXFP8/DeepGEMM reject unwired SiTU; Nemotron ReLU-squared is unsupported.
SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=0uses external FP32 accumulation. Setting it to
1permits nondeterministicBF16 atomic candidates; tuning may still select external reduction. Use native
supports_output_viewto distinguish W4A4's borrowed output from W4A16's owned output.prepared weights into existing buffers. Retain the live layer, workspace owner
and captured addresses; stage runtime scales on every forward. Initialize dummy
weight_scale_2to one while preserving ordinary random-weight initialization.Serialized W4A16 reload is covered; W4A4 reload and online RL reload are outside scope.
Startup tuning and cache ownership
workspaces: powers-of-two decode profiles cover captured, eager and adaptive-MTP
bounds, plus the per-rank prefill ceiling from the shared runtime helper.
All ranks choose the smallest capacity covering synchronized DP counts,
including empty ranks; reject overflow.
Native first-forward AUTO alone would tune the first live input while caching
under workspace capacity. FlashInfer owns namespaced persistence; stage each
profile once across layers and preserve session winners for pooled workspaces.
GPUs do not overwrite each other's records. Same-GPU engines need separate
SGLANG_CACHE_DIRroots. Cache-disabled runs use unique output filenames.Cache naming adds no forward-path work.
SGLANG_FLASHINFER_AUTOTUNE_EXTENDimplicitly. Skipped/disabled tuning retainsthe default native path. Native tuning requires EP=WORLD; draft MegaMoE profiles
and new deterministic-mode behavior are outside scope.
Single-node NVSHMEM defaults
For supported single-node FlashInfer MegaMoE formats with no speculative decoding
or draft A2A
none, set these defaults before NVSHMEM initialization, preservingexplicit environment overrides:
These select the direct-P2P path and avoid the uninitialized proxy global-exit
state corrected by NVIDIA/nvshmem commit
af9e2391a849.NVIDIA/nvshmem#110 is closed; its
description and discussion provide additional background.
All EP GPUs must support peer memory access. Proxy-based global exit/timeout
polling and remote IB/UCX transport are disabled; NCCL/UID bootstrap networking
is unaffected. Multi-node jobs and drafts using other A2A backends receive no
automatic defaults. No dedicated environment-variable behavior tests are added.
Example on four Blackwell GPUs with FlashInfer 0.7.1:
Accuracy Tests
SGLang
38114937747fff49d749c3ee6a7e2f240e34530awith official FlashInfer0.7.1 wheels (release metadata
0b46115d238d29ea2e92c6b9b2cb077d40fbc352).All nine wheel URLs/hashes match release and install metadata; 363 non-FI package
versions/locations and six SGLang native hashes match the image. Wheel installation
uses no source build; CuTe runtime JIT still occurs at first use.
Environment: C2
hu-pdx-111, eight B300 SXM6 AC GPUs (SM103; 275040 MiB each),four GPUs per native/model run; one visible GPU for units. Image
lmsysorg/sglang:nightly-dev-cu13-20261010-961bcf48, digestsha256:e03dd310dc72e64ac709ceb244a79f20fd5e05448c811ab33cc777d8040061ff.Driver 590.48.01; Python 3.12.15; Torch 2.14.1+cu130/CUDA 13.0; NVCC 13.0.88;
Triton 3.8.0; CuTe DSL/libraries 4.8.0; Transformers 5.19.0;
sglang-kernel0.4.9; NVSHMEM 3.4.5/Python 0.4.0; NCCL 2.30.7.FlashInfer imports installed site-packages; SGLang imports the clean pinned checkout.
Reproduce the released wheel environment
Python/cubin are 0.7.1; the aggregate and six architecture JIT-provider wheels
are 0.7.1+cu130. Setup archive SHA256:
7ed1016101cf70523c93eed2d8f475157854b7b32d9456ba91c37681e598d35f.Registered units: actual runtime imports; unit cases mock kernel calls.
Validation helper: smoke_autotune_gpu.py
Native EP4: two layers, hidden/intermediate 256, experts 8/top-k 4; real kernels,
tuning, communication and graph replay with synthetic scheduler counts. Checks:
capacities
[1,2,4,8,16], two startup contexts, rank-cache equality, pooled winners,counts
[0,1,1,2]/[0,1,2,7]/[8,8,8,8]/[0,4,9,16], runtime scale updates andpool refs 2→1→0.
silumeans SwiGLU; SiTU alpha 4/clamp 25 map to beta/linear-beta.Observed AUTO calls: 5/rank cold, 0 warm; no candidate was replaced/disabled.
Harness SHA256:
56e343a213fa155ac00a8a6567d20ca213fd61ffac5bc65314ace58ad0e3c6c4.All four launches exited 0. Complete native fields shared by all 16 reports:
{"passed":true,"precision":"w4a16","capacities":[1,2,4,8,16],"bootstrap_capacity":4,"decode_num_tokens":16,"prefill_num_tokens":8,"startup_contexts":2,"checked_counts":[[0,1,1,2],[0,1,2,7],[8,8,8,8],[0,4,9,16]],"graph_replays":12,"bit_exact":true,"pool_refcounts":[2,1,0]}Model checks: BF16 expert activations/EP transport, TP=attention-DP=EP=4, MegaMoE runner/A2A, W4A16 and
external FC2 reduction. Online converts all routed experts with the unchanged
4over6 recipe and dense source semantics; GSM8K: 200 examples, 128 threads,
completion API, 512-token cap,
score > 0.90. Reload covers both cache-flush modesand memory release/resume; exact text/token IDs and logprob difference ≤
1e-4.The wrapper changes only snapshot paths/ports, keeps registered flags/assertions,
and writes the dummy fixture manifest.
Online runtime: BF16 KV on all four ranks (
autoin server-info); effectivemax-running-requests=3839/DP. Observed 14 tuned capacities:
1,2,4,8,16,32,64,128,256,512,1024,2048,3839,4096.Raw online footer (exit 0):
Reload runtime (startup logs): FP8 E4M3 KV, max-running-requests=4096/DP;
13 power-of-two capacities 1–4096, 52 decode graph sizes through 512. Both
flush modes passed release/resume, exact text/token and 1e-4 logprob checks.
Raw reload footer (exit 0):
Validation helper: run_dummy_megamoe.py
Equivalent reproduction with the same server arguments and assertions; the tested helper additionally records fixture and worker metadata.
Dummy loader: pinned NVIDIA config/tokenizer, only two layers overridden;
real hidden/expert geometry. BF16 combine, external FC2 reduction, per-token
activation 0, TP=attention-DP=EP=4; configured context 2048/prefill 256/requests 16/decode graph 8,
prefill graphs off/memory 0.4. Resolved per-DP prefill 64/requests 4; captured
[1,2,4], FP8 E4M3 KV. Each DP rank generates eight finite-logprob tokensfrom IDs
[1,2,3,4], temperature 0/ignore-EOS. Helper records responses and workermetadata, then stops its server. SHA256:
b61ff2b30562bec77ff86e8826b5d72d621c8c151e52a3ad5d543afe1ac1a407.Scope: native references are independent fixed-default sessions, not a
PyTorch oracle. SiTU uses a small EP4 geometry, not full K3 E2E. Dummy startup
is not accuracy. No W4A4 model-accuracy/reload, online RL reload, Nemotron or
concurrent-engine cache claim. Correctness timings are not serving performance measurements.
Speed Tests and Profiling
End-to-end serving performance has not been measured on the current head.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #38022492097
Latest PR Test (Extra): ❌ Run #38022491936
Latest PR Test (AMD ROCm 10): ❌ Run #38022492254