fix(moe): release W4A16 CuTe DSL scratch workspaces - #4558
Merged
IwakuraRein merged 2 commits intoAug 18, 2026
Merged
IwakuraRein merged 2 commits into
IwakuraRein merged 2 commits into
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
Included review availability: Your plan includes up to 8 reviews per rolling hour; 6 remain after this review. 📝 WalkthroughWalkthroughThe W4A16 MoE launcher now allocates workspace buffers per invocation. Its workspace cache parameter and capacity state were removed. The tuner no longer initializes or passes workspace-cache state. ChangesW4A16 workspace lifecycle
Estimated code review effort: 3 (Moderate) | ~15–30 minutes Merge Risk: ⚪ Minimal · up to The change makes W4A16 scratch workspaces invocation-local, preventing retained GPU buffers and related out-of-memory failures without changing the public behavior or numerical contract; no actionable merge-blocking risk remains after normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
5 tasks
14 tasks
zianglih
marked this pull request as ready for review
August 17, 2026 18:31
zianglih
requested review from
Aneureka,
IwakuraRein,
StudyingShao,
aleozlx,
feih-nv,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
August 17, 2026 18:31
Contributor
Author
|
@coderabbitai full review |
Contributor
✅ Action performedFull review finished. |
aleozlx
approved these changes
Aug 17, 2026
Member
|
/bot run tests/moe |
Collaborator
Member
|
public CI looks clean adding 0.6.18 label. good to merge & cherry pick once bot run is good |
Collaborator
|
[SUCCESS] Pipeline #63176766: 16/16 executed test jobs passed |
5 tasks
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
## 📌 Description @HumansAnd Fixes W4A16 CuTe DSL MoE scratch-workspace retention exposed by the companion SGLang Nemotron-3 integration. - **Root cause:** each `CuteDslFusedMoEW4A16Runner` retained its sort, permuted-hidden, and intermediate BF16 tensors in `_workspace_cache`. The key includes the route tile, so an autotuned runner could retain six large arenas; model serving then retained a selected arena through every layer-owned wrapper. These are live tensor references, so `torch.cuda.empty_cache()` cannot release them. - **Observed geometry:** at `M=16384`, `top_k=22`, 128 local experts, hidden 1024, intermediate 2688, and route tile 128, one runner retains 735.75 MiB of hidden scratch plus 1,931.34 MiB of intermediate scratch and sort buffers. The reproduced 1.89 GiB OOM is the intermediate allocation. - **Fix:** make W4A16 scratch invocation-local, matching the established W4A4 wrapper path. PyTorch's caching allocator reuses eager allocations after the current-stream work completes, while CUDA graph-private pools own captured storage and replay addresses. - **API scope:** no shared/module-global arena, no wrapper sharing, and no new stream or graph serialization contract. The top-level wrapper, kernel inputs, tactics, outputs, and numerical contract are unchanged; the lower-level runner-only `launch_w4a16_moe(..., workspace_cache=...)` keyword is removed. - **Validation history:** two temporary behavioral tests at [`678867c3`](zianglih@678867c) verified eager workspace lifetime and two-call CUDA-graph replay. [`c5ea45e4`](zianglih@c5ea45e) removes them from the final diff for a smaller PR while preserving them in branch history. - **Consumer integration:** [sgl-project/sglang#35120](sgl-project/sglang#35120) depends on this fix and validates online-quantized Nemotron W4A16 accuracy plus serialized Qwen3 W4A16 disk-reload invariance. ## 🔍 Related Issues - W4A16 support was introduced in [flashinfer-ai#4048](flashinfer-ai#4048); Nemotron SiTU support was added in [flashinfer-ai#4394](flashinfer-ai#4394). - Required by [sgl-project/sglang#35120](sgl-project/sglang#35120). ## Validation Environment: - Image: `lmsysorg/sglang:nightly-dev-cu13-20260817-d91c3682` - Image digest: `sha256:fa8774dd128600a09fd6d46670b06fb69a55dac8a3881e50ccf0916a45eb39af` - Hardware: 8x NVIDIA B300 (SM103), driver `590.48.01` - PyTorch: `2.13.0+cu130`; CUDA: `13.0`; NCCL: `2.29.7` - FlashInfer base: `nightly-v0.6.18-20260817` / `e77a4a0d276367895c3b50a642fd8f326c03fb72` - FlashInfer validation commit: `678867c392be6d76f8004e8533c379eb0c222e0f`; final head: `c5ea45e4422e4c8be8059d79b27581971891132a` - CuTe DSL: `nvidia-cutlass-dsl==4.7.0`; Nsight Systems: `2026.4.1` ### Workspace lifetime and CUDA graph validation The following ran at implementation commit `678867c3`; the final head changes only by removing these temporary tests. ```bash CUDA_VISIBLE_DEVICES=0 \ PYTHONPATH=/hai-workspace/flashinfer-runtime-w4a16-fix:/hai-workspace/flashinfer-w4a16-workspace \ python3 -m pytest -vv -s \ /hai-workspace/flashinfer-w4a16-workspace/tests/moe/test_cute_dsl_fused_moe.py::TestCuteDslMoeW4A16::test_wrapper_releases_eager_workspace_with_live_wrappers \ /hai-workspace/flashinfer-w4a16-workspace/tests/moe/test_cute_dsl_fused_moe.py::TestCuteDslMoeW4A16::test_wrapper_cuda_graph_owns_invocation_workspace ``` ```text # Targeted workspace-lifetime and CUDA-graph regression test_wrapper_releases_eager_workspace_with_live_wrappers PASSED test_wrapper_cuda_graph_owns_invocation_workspace PASSED 2 passed, 70 warnings in 918.73s # Full TestCuteDslMoeW4A16 class (14 existing + 2 temporary tests) 16 passed, 598 warnings in 1133.16s ``` All reported warnings are CuTe DSL or dependency deprecation warnings. ### SGLang Nemotron-3 end-to-end The companion SGLang PR validates online expert-weight quantization with the default 16,384-token prefill-graph capacity and `--mem-fraction-static 0.7`. ```text # Online FP8 checkpoint with online NVFP4 expert-weight quantization Capture target prefill CUDA graph end. elapsed=205.47 s, mem usage=4.74 GB, avail mem=56.87 GB. max_total_num_tokens=30259081, chunked_prefill_size=16384, max_prefill_tokens=16384, available_gpu_mem=54.84 GB Total latency: 127.126 s Score: 0.975 Output throughput: 631.254 token/s avg_spec_accept_length=3.1296052121306137 Ran 1 test in 1890.684s OK ``` The serialized ModelOpt Qwen3 test also completed two same-checkpoint disk reloads across weight/CUDA-graph release and resume while preserving deterministic decode text, token IDs, and logprobs: ```text # Serialized nvidia/Qwen3-30B-A3B-NVFP4 W4A16 disk reload Ran 1 test in 1205.499s OK ``` ### PR flashinfer-ai#4048 distributed benchmark compatibility This compatibility rerun executed at final PR head `c5ea45e4422e4c8be8059d79b27581971891132a`. `benchmarks/bench_cute_dsl_moe_distributed.py` is byte-identical between flashinfer-ai#4048's merged head and this head (`sha256:f61bd31757758bc661ddc94551d04c8a762a27612e0fad10c74c1c3240cc87fa`). The imported runtime copy is also byte-identical to the final head for both changed source files: - `flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py`: `sha256:cd11e0c1c37653b0a58403de04fa31ea20ae75a6ebab1b8330e7091e32d729f7` - `flashinfer/fused_moe/cute_dsl/tuner.py`: `sha256:c2dbbad8f2643f1fcd7f0677ec759d150910f6c44418f261cebd5dcba4c2fd39` Each of the three complete runs uses 10 warmups and 100 cold-L2, maximum-rank CUDA-event samples per cell. The HAI pod additionally needs `PET_LOCAL_ADDR=127.0.0.1` because its hostname is not resolvable in pod-local DNS. `--allreduce-backend auto` selected FlashInfer MNNVL, confirmed by a separate one-token `--verbose` diagnostic in the same environment. ```bash set -euo pipefail export FLASHINFER_SOURCE=/hai-workspace/flashinfer-w4a16-workspace export FLASHINFER_RUNTIME=/hai-workspace/flashinfer-runtime-w4a16-fix export RESULTS_DIR=/hai-workspace/sglang-cutedsl-w4a16-results/pr4048-distributed-c5ea-20260817b export PYTHONPATH="${FLASHINFER_RUNTIME}:${FLASHINFER_SOURCE}" export PET_LOCAL_ADDR=127.0.0.1 export FLASHINFER_DISABLE_VERSION_CHECK=1 export FLASHINFER_NVFP4_4OVER6=1 export FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=1 export FLASHINFER_NVFP4_4OVER6_ERR_MODE=MSE export FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 export FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1 mkdir -p "${RESULTS_DIR}" for run in 1 2 3; do torchrun --standalone --nproc-per-node=8 \ "${FLASHINFER_SOURCE}/benchmarks/bench_cute_dsl_moe_distributed.py" \ --warmup 10 --iters 100 \ --use-per-token-activation --no-fused-finalize \ 2>&1 | tee "${RESULTS_DIR}/end-to-end-run-${run}.log" done ``` All three runs emitted the expected 52 rows: 13 token counts x EP8/TP8 x W4A4/W4A16. <details> <summary>Run 1: all 52 raw DISTRIBUTED_CSV rows</summary> ```text DISTRIBUTED_CSV,ep,w4a4,1,8,1.168512 DISTRIBUTED_CSV,ep,w4a16,1,8,0.944432 DISTRIBUTED_CSV,ep,w4a4,2,8,1.226576 DISTRIBUTED_CSV,ep,w4a16,2,8,0.957840 DISTRIBUTED_CSV,ep,w4a4,4,8,1.222032 DISTRIBUTED_CSV,ep,w4a16,4,8,0.916528 DISTRIBUTED_CSV,ep,w4a4,8,8,1.184704 DISTRIBUTED_CSV,ep,w4a16,8,8,0.963344 DISTRIBUTED_CSV,ep,w4a4,16,8,1.224800 DISTRIBUTED_CSV,ep,w4a16,16,8,0.951408 DISTRIBUTED_CSV,ep,w4a4,32,8,1.234592 DISTRIBUTED_CSV,ep,w4a16,32,8,0.967936 DISTRIBUTED_CSV,ep,w4a4,64,8,1.239088 DISTRIBUTED_CSV,ep,w4a16,64,8,0.958080 DISTRIBUTED_CSV,ep,w4a4,128,8,1.227296 DISTRIBUTED_CSV,ep,w4a16,128,8,0.943760 DISTRIBUTED_CSV,ep,w4a4,256,8,1.228256 DISTRIBUTED_CSV,ep,w4a16,256,8,0.955056 DISTRIBUTED_CSV,ep,w4a4,512,8,1.220816 DISTRIBUTED_CSV,ep,w4a16,512,8,0.949472 DISTRIBUTED_CSV,ep,w4a4,1024,8,1.217200 DISTRIBUTED_CSV,ep,w4a16,1024,8,0.950672 DISTRIBUTED_CSV,ep,w4a4,2048,8,1.226176 DISTRIBUTED_CSV,ep,w4a16,2048,8,1.002928 DISTRIBUTED_CSV,ep,w4a4,4096,8,1.270112 DISTRIBUTED_CSV,ep,w4a16,4096,8,1.163104 DISTRIBUTED_CSV,tp,w4a4,1,8,1.246640 DISTRIBUTED_CSV,tp,w4a16,1,8,0.929936 DISTRIBUTED_CSV,tp,w4a4,2,8,1.193216 DISTRIBUTED_CSV,tp,w4a16,2,8,0.905024 DISTRIBUTED_CSV,tp,w4a4,4,8,1.191968 DISTRIBUTED_CSV,tp,w4a16,4,8,0.914336 DISTRIBUTED_CSV,tp,w4a4,8,8,1.232912 DISTRIBUTED_CSV,tp,w4a16,8,8,1.005168 DISTRIBUTED_CSV,tp,w4a4,16,8,1.275024 DISTRIBUTED_CSV,tp,w4a16,16,8,0.988672 DISTRIBUTED_CSV,tp,w4a4,32,8,1.302352 DISTRIBUTED_CSV,tp,w4a16,32,8,0.991168 DISTRIBUTED_CSV,tp,w4a4,64,8,1.287776 DISTRIBUTED_CSV,tp,w4a16,64,8,1.007712 DISTRIBUTED_CSV,tp,w4a4,128,8,1.288544 DISTRIBUTED_CSV,tp,w4a16,128,8,0.976032 DISTRIBUTED_CSV,tp,w4a4,256,8,1.294512 DISTRIBUTED_CSV,tp,w4a16,256,8,1.014432 DISTRIBUTED_CSV,tp,w4a4,512,8,1.325776 DISTRIBUTED_CSV,tp,w4a16,512,8,1.021552 DISTRIBUTED_CSV,tp,w4a4,1024,8,1.333536 DISTRIBUTED_CSV,tp,w4a16,1024,8,1.087632 DISTRIBUTED_CSV,tp,w4a4,2048,8,1.386016 DISTRIBUTED_CSV,tp,w4a16,2048,8,1.226896 DISTRIBUTED_CSV,tp,w4a4,4096,8,1.534096 DISTRIBUTED_CSV,tp,w4a16,4096,8,1.496576 ``` </details> <details> <summary>Run 2: all 52 raw DISTRIBUTED_CSV rows</summary> ```text DISTRIBUTED_CSV,ep,w4a4,1,8,1.198144 DISTRIBUTED_CSV,ep,w4a16,1,8,0.936384 DISTRIBUTED_CSV,ep,w4a4,2,8,1.208656 DISTRIBUTED_CSV,ep,w4a16,2,8,0.914576 DISTRIBUTED_CSV,ep,w4a4,4,8,1.227664 DISTRIBUTED_CSV,ep,w4a16,4,8,0.951744 DISTRIBUTED_CSV,ep,w4a4,8,8,1.208256 DISTRIBUTED_CSV,ep,w4a16,8,8,1.068048 DISTRIBUTED_CSV,ep,w4a4,16,8,1.231184 DISTRIBUTED_CSV,ep,w4a16,16,8,0.946000 DISTRIBUTED_CSV,ep,w4a4,32,8,1.234944 DISTRIBUTED_CSV,ep,w4a16,32,8,0.953472 DISTRIBUTED_CSV,ep,w4a4,64,8,1.382608 DISTRIBUTED_CSV,ep,w4a16,64,8,0.955488 DISTRIBUTED_CSV,ep,w4a4,128,8,1.217792 DISTRIBUTED_CSV,ep,w4a16,128,8,0.942240 DISTRIBUTED_CSV,ep,w4a4,256,8,1.223472 DISTRIBUTED_CSV,ep,w4a16,256,8,0.939152 DISTRIBUTED_CSV,ep,w4a4,512,8,1.224704 DISTRIBUTED_CSV,ep,w4a16,512,8,0.953024 DISTRIBUTED_CSV,ep,w4a4,1024,8,1.231840 DISTRIBUTED_CSV,ep,w4a16,1024,8,0.955184 DISTRIBUTED_CSV,ep,w4a4,2048,8,1.230240 DISTRIBUTED_CSV,ep,w4a16,2048,8,0.984656 DISTRIBUTED_CSV,ep,w4a4,4096,8,1.246976 DISTRIBUTED_CSV,ep,w4a16,4096,8,1.152400 DISTRIBUTED_CSV,tp,w4a4,1,8,1.205248 DISTRIBUTED_CSV,tp,w4a16,1,8,0.952688 DISTRIBUTED_CSV,tp,w4a4,2,8,1.221104 DISTRIBUTED_CSV,tp,w4a16,2,8,0.927984 DISTRIBUTED_CSV,tp,w4a4,4,8,1.222752 DISTRIBUTED_CSV,tp,w4a16,4,8,0.930240 DISTRIBUTED_CSV,tp,w4a4,8,8,1.264928 DISTRIBUTED_CSV,tp,w4a16,8,8,1.002752 DISTRIBUTED_CSV,tp,w4a4,16,8,1.280784 DISTRIBUTED_CSV,tp,w4a16,16,8,0.994400 DISTRIBUTED_CSV,tp,w4a4,32,8,1.280864 DISTRIBUTED_CSV,tp,w4a16,32,8,1.002256 DISTRIBUTED_CSV,tp,w4a4,64,8,1.291056 DISTRIBUTED_CSV,tp,w4a16,64,8,1.005040 DISTRIBUTED_CSV,tp,w4a4,128,8,1.281040 DISTRIBUTED_CSV,tp,w4a16,128,8,1.001472 DISTRIBUTED_CSV,tp,w4a4,256,8,1.286368 DISTRIBUTED_CSV,tp,w4a16,256,8,0.997664 DISTRIBUTED_CSV,tp,w4a4,512,8,1.297824 DISTRIBUTED_CSV,tp,w4a16,512,8,1.012064 DISTRIBUTED_CSV,tp,w4a4,1024,8,1.319184 DISTRIBUTED_CSV,tp,w4a16,1024,8,1.068624 DISTRIBUTED_CSV,tp,w4a4,2048,8,1.358832 DISTRIBUTED_CSV,tp,w4a16,2048,8,1.234288 DISTRIBUTED_CSV,tp,w4a4,4096,8,1.534160 DISTRIBUTED_CSV,tp,w4a16,4096,8,1.490592 ``` </details> <details> <summary>Run 3: all 52 raw DISTRIBUTED_CSV rows</summary> ```text DISTRIBUTED_CSV,ep,w4a4,1,8,1.154336 DISTRIBUTED_CSV,ep,w4a16,1,8,0.908048 DISTRIBUTED_CSV,ep,w4a4,2,8,1.202480 DISTRIBUTED_CSV,ep,w4a16,2,8,0.942320 DISTRIBUTED_CSV,ep,w4a4,4,8,1.217984 DISTRIBUTED_CSV,ep,w4a16,4,8,0.925856 DISTRIBUTED_CSV,ep,w4a4,8,8,1.219968 DISTRIBUTED_CSV,ep,w4a16,8,8,0.925536 DISTRIBUTED_CSV,ep,w4a4,16,8,1.211344 DISTRIBUTED_CSV,ep,w4a16,16,8,0.944256 DISTRIBUTED_CSV,ep,w4a4,32,8,1.227168 DISTRIBUTED_CSV,ep,w4a16,32,8,0.924864 DISTRIBUTED_CSV,ep,w4a4,64,8,1.209744 DISTRIBUTED_CSV,ep,w4a16,64,8,0.942160 DISTRIBUTED_CSV,ep,w4a4,128,8,1.222128 DISTRIBUTED_CSV,ep,w4a16,128,8,0.950096 DISTRIBUTED_CSV,ep,w4a4,256,8,1.228112 DISTRIBUTED_CSV,ep,w4a16,256,8,0.955936 DISTRIBUTED_CSV,ep,w4a4,512,8,1.238400 DISTRIBUTED_CSV,ep,w4a16,512,8,0.946848 DISTRIBUTED_CSV,ep,w4a4,1024,8,1.222480 DISTRIBUTED_CSV,ep,w4a16,1024,8,0.950592 DISTRIBUTED_CSV,ep,w4a4,2048,8,1.247856 DISTRIBUTED_CSV,ep,w4a16,2048,8,0.995744 DISTRIBUTED_CSV,ep,w4a4,4096,8,1.271680 DISTRIBUTED_CSV,ep,w4a16,4096,8,1.173120 DISTRIBUTED_CSV,tp,w4a4,1,8,1.248944 DISTRIBUTED_CSV,tp,w4a16,1,8,0.931824 DISTRIBUTED_CSV,tp,w4a4,2,8,1.218480 DISTRIBUTED_CSV,tp,w4a16,2,8,0.910384 DISTRIBUTED_CSV,tp,w4a4,4,8,1.222864 DISTRIBUTED_CSV,tp,w4a16,4,8,0.919216 DISTRIBUTED_CSV,tp,w4a4,8,8,1.287648 DISTRIBUTED_CSV,tp,w4a16,8,8,0.991392 DISTRIBUTED_CSV,tp,w4a4,16,8,1.270656 DISTRIBUTED_CSV,tp,w4a16,16,8,0.994000 DISTRIBUTED_CSV,tp,w4a4,32,8,1.291312 DISTRIBUTED_CSV,tp,w4a16,32,8,0.996016 DISTRIBUTED_CSV,tp,w4a4,64,8,1.303648 DISTRIBUTED_CSV,tp,w4a16,64,8,0.994432 DISTRIBUTED_CSV,tp,w4a4,128,8,1.294864 DISTRIBUTED_CSV,tp,w4a16,128,8,1.004208 DISTRIBUTED_CSV,tp,w4a4,256,8,1.309296 DISTRIBUTED_CSV,tp,w4a16,256,8,1.022560 DISTRIBUTED_CSV,tp,w4a4,512,8,1.310064 DISTRIBUTED_CSV,tp,w4a16,512,8,1.026016 DISTRIBUTED_CSV,tp,w4a4,1024,8,1.332736 DISTRIBUTED_CSV,tp,w4a16,1024,8,1.082720 DISTRIBUTED_CSV,tp,w4a4,2048,8,1.383648 DISTRIBUTED_CSV,tp,w4a16,2048,8,1.229296 DISTRIBUTED_CSV,tp,w4a4,4096,8,1.546464 DISTRIBUTED_CSV,tp,w4a16,4096,8,1.511744 ``` </details> The table below takes the median across the three complete runs for each cell. `speedup = median(W4A4 EP8) / median(strategy)`, so values above 1 favor the strategy. | Tokens | W4A4 EP8 ms | W4A4 TP8 ms | TP speedup | W4A16 EP8 ms | EP speedup | W4A16 TP8 ms | TP speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 1.169 | 1.247 | 0.937x | 0.936 | 1.248x | 0.932 | 1.254x | | 2 | 1.209 | 1.218 | 0.992x | 0.942 | 1.283x | 0.910 | 1.328x | | 4 | 1.222 | 1.223 | 0.999x | 0.926 | 1.320x | 0.919 | 1.329x | | 8 | 1.208 | 1.265 | 0.955x | 0.963 | 1.254x | 1.003 | 1.205x | | 16 | 1.225 | 1.275 | 0.961x | 0.946 | 1.295x | 0.994 | 1.232x | | 32 | 1.235 | 1.291 | 0.956x | 0.953 | 1.295x | 0.996 | 1.240x | | 64 | 1.239 | 1.291 | 0.960x | 0.955 | 1.297x | 1.005 | 1.233x | | 128 | 1.222 | 1.289 | 0.948x | 0.944 | 1.295x | 1.001 | 1.220x | | 256 | 1.228 | 1.295 | 0.949x | 0.955 | 1.286x | 1.014 | 1.211x | | 512 | 1.225 | 1.310 | 0.935x | 0.949 | 1.290x | 1.022 | 1.199x | | 1024 | 1.222 | 1.333 | 0.917x | 0.951 | 1.286x | 1.083 | 1.129x | | 2048 | 1.230 | 1.384 | 0.889x | 0.996 | 1.235x | 1.229 | 1.001x | | 4096 | 1.270 | 1.534 | 0.828x | 1.163 | 1.092x | 1.497 | 0.849x | Geometric-mean speedups versus W4A4 EP8 are `0.940x` for W4A4 TP8, `1.266x` for W4A16 EP8, and `1.179x` for W4A16 TP8. Using the same exported environment above, the Nsight Systems driver launches its own eight-rank workers for all eight cases: two token counts x EP8/TP8 x W4A4/W4A16. ```bash python3 "${FLASHINFER_SOURCE}/benchmarks/bench_cute_dsl_moe_distributed.py" \ --mode profile_nsys --warmup 10 --profile-iters 20 \ --use-per-token-activation --no-fused-finalize \ --nsys-output-dir "${RESULTS_DIR}/nsys" \ 2>&1 | tee "${RESULTS_DIR}/nsys.log" ``` <details> <summary>32-token cases: four raw Nsight Systems stage tables</summary> ```text Nsight Systems kernel breakdown: EP8 W4A4, 32 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel routing | 1 | 0.35% | 0.559 | 160 | 3.491 | deepseek_v3_topk_kernel dispatch | 1 | 0.15% | 0.240 | 160 | 1.502 | moeA2APrepareDispatchKernel dispatch | 2 | 19.77% | 31.380 | 160 | 196.127 | moeA2ADispatchKernel dispatch | 3 | 11.53% | 18.301 | 160 | 114.379 | moeA2ASanitizeExpertIdsKernel activation prep/quant | 1 | 0.11% | 0.169 | 160 | 1.059 | vectorized_elementwise_kernel activation prep/quant | 2 | 0.47% | 0.750 | 160 | 4.688 | kernel_cutlass_kernel_flashinf_1... local MoE | 1 | 0.39% | 0.619 | 160 | 3.871 | routingIndicesClusterKernel local MoE | 2 | 7.29% | 11.573 | 160 | 72.334 | kernel_cutlass_kernel_flashinf_2... local MoE | 3 | 0.12% | 0.197 | 160 | 1.230 | vectorized_elementwise_kernel local MoE | 4 | 0.81% | 1.282 | 160 | 8.013 | kernel_cutlass_kernel_flashinf_3... local MoE | 5 | 3.28% | 5.204 | 160 | 32.524 | kernel_cutlass_kernel_flashinf_4... local MoE | 6 | 0.88% | 1.400 | 160 | 8.750 | moeUnpermuteKernel combine | 1 | 0.27% | 0.429 | 160 | 2.682 | moeA2APrepareCombineKernel combine | 2 | 54.58% | 86.649 | 160 | 541.557 | moeA2ACombineKernel Nsight Systems kernel breakdown: EP8 W4A16, 32 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel routing | 1 | 0.66% | 0.557 | 160 | 3.483 | deepseek_v3_topk_kernel dispatch | 1 | 0.29% | 0.240 | 160 | 1.498 | moeA2APrepareDispatchKernel dispatch | 2 | 21.84% | 18.310 | 160 | 114.439 | moeA2ADispatchKernel dispatch | 3 | 4.86% | 4.071 | 160 | 25.444 | moeA2ASanitizeExpertIdsKernel local MoE | 1 | 0.76% | 0.639 | 160 | 3.995 | routingIndicesClusterKernel local MoE | 2 | 0.78% | 0.656 | 160 | 4.099 | moePermuteKernel local MoE | 3 | 12.99% | 10.896 | 160 | 68.099 | kernel_cutlass_kernel_flashinf_1... local MoE | 4 | 7.05% | 5.910 | 160 | 36.940 | kernel_cutlass_kernel_flashinf_1... local MoE | 5 | 1.34% | 1.125 | 160 | 7.032 | moeUnpermuteKernel combine | 1 | 0.47% | 0.390 | 160 | 2.440 | moeA2APrepareCombineKernel combine | 2 | 48.96% | 41.055 | 160 | 256.594 | moeA2ACombineKernel Nsight Systems kernel breakdown: TP8 W4A4, 32 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel all-gather | 1 | 13.30% | 35.050 | 160 | 219.062 | ncclDevKernel_AllGather_RING_LL routing | 1 | 0.21% | 0.566 | 160 | 3.535 | deepseek_v3_topk_kernel activation prep/quant | 1 | 0.06% | 0.153 | 160 | 0.954 | vectorized_elementwise_kernel activation prep/quant | 2 | 0.27% | 0.706 | 160 | 4.411 | kernel_cutlass_kernel_flashinf_1... local MoE | 1 | 0.23% | 0.600 | 160 | 3.747 | routingIndicesClusterKernel local MoE | 2 | 4.46% | 11.747 | 160 | 73.420 | kernel_cutlass_kernel_flashinf_2... local MoE | 3 | 0.08% | 0.201 | 160 | 1.257 | vectorized_elementwise_kernel local MoE | 4 | 2.18% | 5.746 | 160 | 35.910 | kernel_cutlass_kernel_flashinf_3... local MoE | 5 | 2.35% | 6.188 | 160 | 38.673 | kernel_cutlass_kernel_flashinf_4... local MoE | 6 | 1.09% | 2.865 | 160 | 17.906 | moeUnpermuteKernel all-reduce | 1 | 75.78% | 199.723 | 160 | 1248.270 | twoshotAllreduceKernel Nsight Systems kernel breakdown: TP8 W4A16, 32 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel all-gather | 1 | 18.23% | 17.083 | 160 | 106.768 | ncclDevKernel_AllGather_RING_LL routing | 1 | 0.61% | 0.568 | 160 | 3.547 | deepseek_v3_topk_kernel local MoE | 1 | 0.63% | 0.591 | 160 | 3.692 | routingIndicesClusterKernel local MoE | 2 | 0.87% | 0.815 | 160 | 5.093 | moePermuteKernel local MoE | 3 | 11.64% | 10.909 | 160 | 68.178 | kernel_cutlass_kernel_flashinf_1... local MoE | 4 | 7.32% | 6.863 | 160 | 42.897 | kernel_cutlass_kernel_flashinf_1... local MoE | 5 | 2.21% | 2.067 | 160 | 12.920 | moeUnpermuteKernel all-reduce | 1 | 58.50% | 54.823 | 160 | 342.642 | twoshotAllreduceKernel ``` </details> <details> <summary>4,096-token cases: four raw Nsight Systems stage tables</summary> ```text Nsight Systems kernel breakdown: EP8 W4A4, 4096 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel routing | 1 | 0.39% | 0.663 | 160 | 4.143 | deepseek_v3_topk_kernel dispatch | 1 | 0.15% | 0.247 | 160 | 1.546 | moeA2APrepareDispatchKernel dispatch | 2 | 18.26% | 30.738 | 160 | 192.110 | moeA2ADispatchKernel dispatch | 3 | 9.38% | 15.786 | 160 | 98.664 | moeA2ASanitizeExpertIdsKernel activation prep/quant | 1 | 0.15% | 0.259 | 160 | 1.620 | vectorized_elementwise_kernel activation prep/quant | 2 | 2.04% | 3.435 | 160 | 21.470 | kernel_cutlass_kernel_flashinf_1... local MoE | 1 | 0.63% | 1.060 | 160 | 6.623 | routingIndicesClusterKernel local MoE | 2 | 10.34% | 17.407 | 160 | 108.791 | kernel_cutlass_kernel_flashinf_2... local MoE | 3 | 0.23% | 0.385 | 160 | 2.405 | vectorized_elementwise_kernel local MoE | 4 | 5.91% | 9.940 | 160 | 62.123 | kernel_cutlass_kernel_flashinf_3... local MoE | 5 | 5.83% | 9.808 | 160 | 61.299 | kernel_cutlass_kernel_flashinf_4... local MoE | 6 | 3.10% | 5.219 | 160 | 32.616 | moeUnpermuteKernel combine | 1 | 0.93% | 1.571 | 160 | 9.817 | moeA2APrepareCombineKernel combine | 2 | 42.66% | 71.811 | 160 | 448.816 | moeA2ACombineKernel Nsight Systems kernel breakdown: EP8 W4A16, 4096 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel routing | 1 | 0.45% | 0.658 | 160 | 4.112 | deepseek_v3_topk_kernel dispatch | 1 | 0.16% | 0.240 | 160 | 1.500 | moeA2APrepareDispatchKernel dispatch | 2 | 13.16% | 19.373 | 160 | 121.083 | moeA2ADispatchKernel dispatch | 3 | 2.70% | 3.982 | 160 | 24.887 | moeA2ASanitizeExpertIdsKernel local MoE | 1 | 0.71% | 1.039 | 160 | 6.494 | routingIndicesClusterKernel local MoE | 2 | 2.20% | 3.241 | 160 | 20.257 | moePermuteKernel local MoE | 3 | 22.48% | 33.100 | 160 | 206.873 | kernel_cutlass_kernel_flashinf_1... local MoE | 4 | 14.91% | 21.957 | 160 | 137.229 | kernel_cutlass_kernel_flashinf_1... local MoE | 5 | 3.47% | 5.106 | 160 | 31.915 | moeUnpermuteKernel combine | 1 | 1.44% | 2.125 | 160 | 13.279 | moeA2APrepareCombineKernel combine | 2 | 38.32% | 56.418 | 160 | 352.613 | moeA2ACombineKernel Nsight Systems kernel breakdown: TP8 W4A4, 4096 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel all-gather | 1 | 20.71% | 49.212 | 160 | 307.574 | ncclDevKernel_AllGather_RING_LL routing | 1 | 0.68% | 1.627 | 160 | 10.170 | deepseek_v3_topk_kernel activation prep/quant | 1 | 0.10% | 0.245 | 160 | 1.534 | vectorized_elementwise_kernel activation prep/quant | 2 | 2.01% | 4.774 | 160 | 29.837 | kernel_cutlass_kernel_flashinf_1... local MoE | 1 | 0.53% | 1.258 | 160 | 7.864 | routingIndicesClusterKernel local MoE | 2 | 8.26% | 19.630 | 160 | 122.688 | kernel_cutlass_kernel_flashinf_2... local MoE | 3 | 0.09% | 0.215 | 160 | 1.345 | vectorized_elementwise_kernel local MoE | 4 | 4.90% | 11.640 | 160 | 72.748 | kernel_cutlass_kernel_flashinf_3... local MoE | 5 | 9.40% | 22.334 | 160 | 139.585 | kernel_cutlass_kernel_flashinf_4... local MoE | 6 | 7.93% | 18.841 | 160 | 117.757 | moeUnpermuteKernel all-reduce | 1 | 45.37% | 107.793 | 160 | 673.706 | twoshotAllreduceKernel Nsight Systems kernel breakdown: TP8 W4A16, 4096 global tokens stage | launch | GPU time | total (ms) | instances | avg (us) | kernel all-gather | 1 | 14.75% | 31.958 | 160 | 199.739 | ncclDevKernel_AllGather_RING_LL routing | 1 | 0.75% | 1.632 | 160 | 10.197 | deepseek_v3_topk_kernel local MoE | 1 | 0.57% | 1.243 | 160 | 7.769 | routingIndicesClusterKernel local MoE | 2 | 9.78% | 21.189 | 160 | 132.433 | moePermuteKernel local MoE | 3 | 16.70% | 36.189 | 160 | 226.178 | kernel_cutlass_kernel_flashinf_1... local MoE | 4 | 18.50% | 40.095 | 160 | 250.596 | kernel_cutlass_kernel_flashinf_1... local MoE | 5 | 8.52% | 18.467 | 160 | 115.418 | moeUnpermuteKernel all-reduce | 1 | 30.42% | 65.919 | 160 | 411.993 | twoshotAllreduceKernel ``` </details> The raw `total (ms)` columns above aggregate all eight ranks and 20 captured iterations. Dividing each case's sum by `8 * 20 = 160` gives: | Tokens | Mode | W4A4 staged ms | W4A16 staged ms | W4A4 / W4A16 | |---:|:---:|---:|---:|---:| | 32 | EP8 | 0.992 | 0.524 | 1.893x | | 32 | TP8 | 1.647 | 0.586 | 2.812x | | 4096 | EP8 | 1.052 | 0.920 | 1.143x | | 4096 | TP8 | 1.485 | 1.354 | 1.096x | Limitations: - This is a current-head compatibility run, not a base-vs-head performance A/B. It does not by itself prove that invocation-local allocation has zero latency overhead. - It ran on B300/SM103 with CUDA 13.0; flashinfer-ai#4048's published tables used B200 with different CUDA/runtime revisions, so absolute timings are not directly comparable. - Two end-to-end cells had isolated run ranges around 14% (EP8 W4A16 at 8 tokens and EP8 W4A4 at 64 tokens); use the reported three-run medians and avoid sub-percent conclusions. - Nsight staged sums are diagnostic kernel-time aggregates, not wall-clock latency, and do not subtract overlapping kernels. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes Please focus on the workspace lifetime and CUDA-graph ownership contract. This intentionally follows W4A4's invocation-local scratch pattern instead of adding a process-global or cross-wrapper workspace cache. The final diff contains no new test code; the validation-only behavioral tests and their passing results remain reviewable at `678867c3`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved W4A16 mixture-of-experts operation reliability by allocating required working memory for each invocation. * Improved handling of workloads with varying sizes by removing reliance on previously cached capacity. * **Refactor** * Simplified workspace management while preserving existing operation behavior. * Streamlined the W4A16 execution interface for more predictable invocation behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
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
4 participants
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.
📌 Description
@HumansAnd
Fixes W4A16 CuTe DSL MoE scratch-workspace retention exposed by the companion SGLang Nemotron-3 integration.
Root cause: each
CuteDslFusedMoEW4A16Runnerretained its sort, permuted-hidden, and intermediate BF16 tensors in_workspace_cache. The key includes the route tile, so an autotuned runner could retain six large arenas; model serving then retained a selected arena through every layer-owned wrapper. These are live tensor references, sotorch.cuda.empty_cache()cannot release them.Observed geometry: at
M=16384,top_k=22, 128 local experts, hidden 1024, intermediate 2688, and route tile 128, one runner retains 735.75 MiB of hidden scratch plus 1,931.34 MiB of intermediate scratch and sort buffers. The reproduced 1.89 GiB OOM is the intermediate allocation.Fix: make W4A16 scratch invocation-local, matching the established W4A4 wrapper path. PyTorch's caching allocator reuses eager allocations after the current-stream work completes, while CUDA graph-private pools own captured storage and replay addresses.
API scope: no shared/module-global arena, no wrapper sharing, and no new stream or graph serialization contract. The top-level wrapper, kernel inputs, tactics, outputs, and numerical contract are unchanged; the lower-level runner-only
launch_w4a16_moe(..., workspace_cache=...)keyword is removed.Validation history: two temporary behavioral tests at
678867c3verified eager workspace lifetime and two-call CUDA-graph replay.c5ea45e4removes them from the final diff for a smaller PR while preserving them in branch history.Consumer integration: sgl-project/sglang#35120 depends on this fix and validates online-quantized Nemotron W4A16 accuracy plus serialized Qwen3 W4A16 disk-reload invariance.
🔍 Related Issues
Validation
Environment:
lmsysorg/sglang:nightly-dev-cu13-20260817-d91c3682sha256:fa8774dd128600a09fd6d46670b06fb69a55dac8a3881e50ccf0916a45eb39af590.48.012.13.0+cu130; CUDA:13.0; NCCL:2.29.7nightly-v0.6.18-20260817/e77a4a0d276367895c3b50a642fd8f326c03fb72678867c392be6d76f8004e8533c379eb0c222e0f; final head:c5ea45e4422e4c8be8059d79b27581971891132anvidia-cutlass-dsl==4.7.0; Nsight Systems:2026.4.1Workspace lifetime and CUDA graph validation
The following ran at implementation commit
678867c3; the final head changes only by removing these temporary tests.All reported warnings are CuTe DSL or dependency deprecation warnings.
SGLang Nemotron-3 end-to-end
The companion SGLang PR validates online expert-weight quantization with the default 16,384-token prefill-graph capacity and
--mem-fraction-static 0.7.The serialized ModelOpt Qwen3 test also completed two same-checkpoint disk reloads across weight/CUDA-graph release and resume while preserving deterministic decode text, token IDs, and logprobs:
PR #4048 distributed benchmark compatibility
This compatibility rerun executed at final PR head
c5ea45e4422e4c8be8059d79b27581971891132a.benchmarks/bench_cute_dsl_moe_distributed.pyis byte-identical between #4048's merged head and this head (sha256:f61bd31757758bc661ddc94551d04c8a762a27612e0fad10c74c1c3240cc87fa). The imported runtime copy is also byte-identical to the final head for both changed source files:flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py:sha256:cd11e0c1c37653b0a58403de04fa31ea20ae75a6ebab1b8330e7091e32d729f7flashinfer/fused_moe/cute_dsl/tuner.py:sha256:c2dbbad8f2643f1fcd7f0677ec759d150910f6c44418f261cebd5dcba4c2fd39Each of the three complete runs uses 10 warmups and 100 cold-L2, maximum-rank CUDA-event samples per cell. The HAI pod additionally needs
PET_LOCAL_ADDR=127.0.0.1because its hostname is not resolvable in pod-local DNS.--allreduce-backend autoselected FlashInfer MNNVL, confirmed by a separate one-token--verbosediagnostic in the same environment.All three runs emitted the expected 52 rows: 13 token counts x EP8/TP8 x W4A4/W4A16.
Run 1: all 52 raw DISTRIBUTED_CSV rows
Run 2: all 52 raw DISTRIBUTED_CSV rows
Run 3: all 52 raw DISTRIBUTED_CSV rows
The table below takes the median across the three complete runs for each cell.
speedup = median(W4A4 EP8) / median(strategy), so values above 1 favor the strategy.Geometric-mean speedups versus W4A4 EP8 are
0.940xfor W4A4 TP8,1.266xfor W4A16 EP8, and1.179xfor W4A16 TP8.Using the same exported environment above, the Nsight Systems driver launches its own eight-rank workers for all eight cases: two token counts x EP8/TP8 x W4A4/W4A16.
32-token cases: four raw Nsight Systems stage tables
4,096-token cases: four raw Nsight Systems stage tables
The raw
total (ms)columns above aggregate all eight ranks and 20 captured iterations. Dividing each case's sum by8 * 20 = 160gives:Limitations:
🚀 Pull Request Checklist
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Please focus on the workspace lifetime and CUDA-graph ownership contract. This intentionally follows W4A4's invocation-local scratch pattern instead of adding a process-global or cross-wrapper workspace cache. The final diff contains no new test code; the validation-only behavioral tests and their passing results remain reviewable at
678867c3.Summary by CodeRabbit