[NVIDIA] Add flashinfer MNNVL backend for allreduce only by wenscarl · Pull Request #30700 · sgl-project/sglang

wenscarl · 2026-07-09T22:08:04Z

Motivation

SGLang already supports fused allreduce via FlashInfer (kARResidualRMSNorm: allreduce + residual + RMSNorm in one kernel) for DeepSeek-V3/V4 models. However, not every allreduce site can use the fused path — RowParallelLinear, attention output fallback, and MoE TP fallback all perform pure (non-fused) allreduce that currently always falls back to NCCL/custom-allreduce even when the FlashInfer mnnvl/trtllm workspace is already initialized and available.

This PR routes those pure allreduce calls through FlashInfer AllReduceFusionPattern.kAllReduce, reusing the existing workspace with no extra memory cost.

Changes

`flashinfer_comm_fusion.py` — new `flashinfer_allreduce()`

Wraps AllReduceFusionPattern.kAllReduce for pure TP allreduce. Returns None on any of:

FlashInfer unavailable or _flashinfer_comm not initialized
Input is not 2D contiguous (e.g. 3D vocab embedding output — graceful NCCL fallback)
Workspace not initialized or buffer too small

`parallel_state.py` — dispatch hook + group tagging

_tag_groups_for_flashinfer_pure_allreduce(): stamps a _fi_workspace_hint attribute ("attn_tp" / "moe_tp" / "moe_ep") on _TP, _ATTN_TP, _MOE_TP, _MOE_EP group coordinators after initialize_model_parallel().
GroupCoordinator.all_reduce(): checks _fi_workspace_hint and dispatches to flashinfer_allreduce() before the existing ca_comm/pymscclpp/NCCL chain. Falls back transparently if flashinfer_allreduce() returns None.

`server_args.py` — new flag

--enable-flashinfer-pure-allreduce (resolvable=True so post-process passes can set it).

`model_runner.py` — wiring

Calls set_flashinfer_pure_all_reduce() and _tag_groups_for_flashinfer_pure_allreduce() alongside the existing set_custom_all_reduce / set_mscclpp_all_reduce setup.

`overrides.py` — auto-enable for DeepSeek-V3/V4

Added DeepseekV4ForCausalLM to _FLASHINFER_ALLREDUCE_FUSION_ARCHS.
New _FLASHINFER_PURE_ALLREDUCE_ARCHS frozenset (DeepseekV3ForCausalLM, DeepseekV32ForCausalLM, DeepseekV4ForCausalLM).
_flashinfer_allreduce_fusion_auto_enable now also sets enable_flashinfer_pure_allreduce=True for archs in _FLASHINFER_PURE_ALLREDUCE_ARCHS — users running DeepSeek-V3/V4 no longer need to pass --flashinfer-allreduce-fusion-backend or --enable-flashinfer-pure-allreduce manually.
_enforce_disable_allreduce_fusion now also clears enable_flashinfer_pure_allreduce when the fused backend is force-disabled.

`test_flashinfer_pure_allreduce.py` — unit tests

5 unit tests using _FakeFlashInferComm covering: correctness (output == input × world_size), 1D/3D shape guard, non-contiguous guard, unavailable flag, uninitialized workspace.

Performance and Accuracy

Benchmarked DeepSeek V4 Flash on Blackwell with sglang.benchmark.one_batch, using TP=4 and sweeping batch sizes.

Command:

python3 -m sglang.benchmark.one_batch \
  --model-path deepseek-ai/DeepSeek-V4-Flash \
  --tp-size 4 \
  --trust-remote-code \
  --batch-size 1 2 4 8 16 32 64 \
  --input-len 512 \
  --output-len 128 \
  --run-name ds_v4_flash_decode_bs_sweep \
  --disable-flashinfer-autotune \
  --result-filename ds_v4_flash_decode_bs_sweep.jsonl \
  --moe-runner-backend flashinfer_mxfp4

BS	Old decode tok/s	PR decode tok/s	Decode delta	Old decode lat	PR decode lat	Old overall tok/s	PR overall tok/s
1	147.51	157.69	+6.9%	6.78 ms	6.34 ms	539.58	624.73
2	285.43	304.95	+6.8%	7.01 ms	6.56 ms	1,014.16	1,185.12
4	535.28	569.52	+6.4%	7.47 ms	7.02 ms	1,959.97	2,269.91
8	1,071.14	1,103.58	+3.0%	7.47 ms	7.25 ms	2,863.78	4,586.34
16	1,966.24	2,016.57	+2.6%	8.14 ms	7.93 ms	8,205.76	8,482.90
32	3,510.57	3,616.18	+3.0%	9.12 ms	8.85 ms	14,871.86	15,267.53
64	6,247.72	6,157.77	-1.4%	10.24 ms	10.39 ms	24,242.08	23,916.24

Accuracy:

python3 benchmark/gsm8k/bench_sglang.py \
  --num-shots 8 \
  --num-questions 1316 \
  --parallel 1316

Without this PR:

Accuracy: 0.958
Invalid: 0.000
Latency: 95.622 s
Output throughput: 1312.239 token/s

With this PR:

Accuracy: 0.958
Invalid: 0.000
Latency: 107.229 s
Output throughput: 1158.934 token/s
``

cc. @Fridge003 

<!-- pr-states:start -->
---
### CI States

Latest PR Test (Base): <!-- slot:pr-test:start -->:white_check_mark: [Run #31440466100](https://github.com/sgl-project/sglang/actions/runs/31440466100)<!-- slot:pr-test:end -->
Latest PR Test (Extra): <!-- slot:pr-test-extra:start -->:white_check_mark: [Run #31547727723](https://github.com/sgl-project/sglang/actions/runs/31547727723)<!-- slot:pr-test-extra:end -->
<!-- pr-states:end -->

gemini-code-assist · 2026-07-09T22:08:08Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

nvpohanh · 2026-07-15T02:35:56Z

@wenscarl Could you resolve the conflicts and set this PR as ready for review? Thanks!

nvpohanh · 2026-07-15T07:07:07Z

@shyeh25 could you share your perf results in this PR comment? thanks

shyeh25 · 2026-07-15T07:36:35Z

https://github.com/SemiAnalysisAI/InferenceX/blob/main/benchmarks/multi_node/srt-slurm-recipes/sglang/deepseek-v4/8k1k/disagg-gb200-1p1d-tp8-tp8-4-c1.yaml

1P1D TP8-TP8 conc1	TPS per GPU	TPS per user
GB200 w/o PR30700	78.77	79.43
GB200 w/ PR30700	97.12	100.81

https://github.com/SemiAnalysisAI/InferenceX/blob/main/benchmarks/multi_node/srt-slurm-recipes/sglang/deepseek-v4/8k1k/disagg-gb200-1p4d-dep8-tp8-10-c64.yaml

1P4D-DEP8-TP8 conc64	TPS per GPU	TPS per user
GB200 w/o PR30700	726.36	62.46
GB200 w/ PR30700	897.05	81.83

gemini-code-assist · 2026-07-16T15:28:46Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

nvpohanh · 2026-07-17T06:57:46Z

SGLang already supports fused allreduce via FlashInfer (kARResidualRMSNorm: allreduce + residual + RMSNorm in one kernel) for DeepSeek-V3/V4 models.

Correction about this: DSV4 has mHC so AR+Norm cannot be fused. so SGLang only supported FlashInfer MNNVL AR for DSV3 but not DSV4. This PR adds that support for DSV4 and can improve DSV4 GB200 TP8 small-conc perf significantly.

nvpohanh · 2026-07-17T12:28:33Z

@wenscarl could you fix the linter failure? thanks!

nvpohanh · 2026-07-20T03:23:08Z

/tag-and-rerun-ci

ormandj · 2026-07-24T15:11:40Z

We tested this PR end to end on SM120 (2x RTX PRO 6000 Blackwell, TP2, DeepSeek-V4-Flash MTP2, CUDA graphs) on top of main 3d91a569ce, with FlashInfer 0.6.14 + flashinfer-ai/flashinfer#3903 + flashinfer-ai/flashinfer#3930. Two findings:

As written, the PR fails closed on SM120: _resolve_backend rejects the arch, and the workspace preflight's multicast-granularity query rejects the non-multicast TRT-LLM workspace. With five consecutive 30 s runs per cell on a settled server: PR [NVIDIA] Add flashinfer MNNVL backend for allreduce only #30700 alone medians 126.25 tok/s (mean 125.75; samples 125.90, 127.66, 130.19, 126.25, 118.75), and with two small companion commits (SM120 in the selector/auto-enable, multicast-only preflight probe removed) it medians 132.85 (mean 134.31; samples 129.55, 131.81, 132.85, 140.27, 137.05) — +5.2% on medians, +6.8% on means. All 87 per-graph NCCL ring all-reduce kernels are replaced 1:1 by one-shot Lamport kernels, 0.772 → 0.394 ms per graph, with GSM8K parity (97.5% vs 97.0%, 200 questions). We opened fix: enable FlashInfer TRT-LLM all-reduce on SM12X #32330 with those commits stacked on this PR; happy to fold them in here instead if you prefer.
This PR's change to _enforce_disable_allreduce_fusion (returning enable_flashinfer_allreduce_only: False alongside the backend reset) breaks test/registered/unit/test_model_overrides.py::test_flashinfer_allreduce_fusion_passes on every architecture — the test still expects the single-key dict {"flashinfer_allreduce_fusion_backend": None}. Reproduced at head 22e2f8b30d merged onto main 3d91a569ce. One-line expectation fix.

nvpohanh · 2026-07-28T08:54:56Z

@wenscarl Could you address the comments and fix the conflicts? thanks!

b8zhong · 2026-07-29T14:40:55Z

/rerun-test test/registered/unit/layers/test_flashinfer_comm_fusion.py

github-actions · 2026-07-29T14:41:27Z

Results for /rerun-test test/registered/unit/layers/test_flashinfer_comm_fusion.py:

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

cd test/ && python3 registered/unit/layers/test_flashinfer_comm_fusion.py

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

cd test/ && python3 registered/unit/layers/test_flashinfer_comm_fusion.py

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

cd test/ && python3 registered/unit/layers/test_flashinfer_comm_fusion.py

github-actions · 2026-08-12T01:09:25Z

Results for /rerun-test test/registered/sampling/test_penalty.py test/registered/radix_cache/test_mamba2_extra_buffer_kl.py test/registered/spec/test_gemma4_dflash_31b_extra.py:

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

cd test/ && python3 registered/sampling/test_penalty.py

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

cd test/ && python3 registered/radix_cache/test_mamba2_extra_buffer_kl.py

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/test_gemma4_dflash_31b_extra.py

…#30700)

sgl-project#30700 made should_fuse_mlp_allreduce_with_next_layer read parallel.moe_ep_size and parallel.moe_tp_size, but the get_parallel mock in the AMD gate test only carries tp_size, so every case raises AttributeError. Set both to 1, which is the dense-TP configuration these tests exercise.

…e test stub The hybrid EP+TP guard added in #30700 reads parallel.moe_ep_size and parallel.moe_tp_size at the top of should_fuse_mlp_allreduce_with_next_layer, but the AMD gate test still stubs get_parallel() with a SimpleNamespace that only carries tp_size, so every gate case raises AttributeError.

The kAllReduce path added in sgl-project#30700 hardcoded fp32_acc=False, so it kept rounding to the input dtype once per rank after the fused path stopped.

…#30700)

wenscarl force-pushed the flashinfer-pure-allreduce branch from 59ab05c to 1ed97b3 Compare July 16, 2026 15:27

wenscarl marked this pull request as ready for review July 16, 2026 15:28

wenscarl requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Ying1123, ch-wan, ispobock, merrymercy and yizhang2077 as code owners July 16, 2026 15:28

wenscarl changed the title ~~Flashinfer pure allreduce~~ [NVIDIA] Add flashinfer MNNVL backend for allreduce only Jul 17, 2026

github-actions Bot added the run-ci label Jul 20, 2026

This was referenced Jul 24, 2026

fix: enable FlashInfer TRT-LLM all-reduce on SM12X #32330

Closed

feat(comm): extend trtllm_allreduce to SM12x and fix lamport buffer pointer packing flashinfer-ai/flashinfer#3903

Merged

b8zhong self-assigned this Jul 29, 2026

b8zhong reviewed Jul 29, 2026

View reviewed changes

Comment thread python/sglang/srt/arg_groups/overrides.py Outdated

b8zhong reviewed Jul 29, 2026

View reviewed changes

Comment thread python/sglang/srt/server_args.py Outdated

vstone-w pushed a commit to ClownBin/sglang that referenced this pull request Aug 12, 2026

[NVIDIA] Add flashinfer MNNVL backend for allreduce only (sgl-project…

6fb574c

…#30700)

amd-bot mentioned this pull request Aug 12, 2026

[CI Monitor] Daily Report - 2026-08-12 bingxche/sglang-ci-bot#164

Open

b8zhong mentioned this pull request Aug 12, 2026

[Feature] Should the trtllm allreduce fusion accumulate in fp32, like the mnnvl backends do? #34603

Closed

2 tasks

saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026

[NVIDIA] Add flashinfer MNNVL backend for allreduce only (sgl-project…

9241196

…#30700)

This was referenced Aug 17, 2026

[CI Monitor] Daily Report - 2026-08-17 bingxche/sglang-ci-bot#169

Open

[CI Monitor] Daily Report - 2026-08-18 bingxche/sglang-ci-bot#170

Open

kangwangamd mentioned this pull request Aug 18, 2026

[AMD] Add moe_ep_size/moe_tp_size to the allreduce-fusion gate test mock #35251

Closed

5 tasks

bingxche mentioned this pull request Aug 18, 2026

[AMD][bugfix] Add moe_ep_size/moe_tp_size to the allreduce-fusion gate test stub #35340

Merged

5 tasks

learning-sketch mentioned this pull request Aug 24, 2026

Sync upstream sglang main and refresh DSV4 framework PR checklist learning-sketch/sglang#8

Draft

This was referenced Aug 24, 2026

[CI Monitor] Daily Report - 2026-08-24 bingxche/sglang-ci-bot#176

Open

[CI Monitor] Daily Report - 2026-08-25 bingxche/sglang-ci-bot#177

Open

[CI Monitor] Daily Report - 2026-08-26 bingxche/sglang-ci-bot#178

Open

lelvisl mentioned this pull request Sep 2, 2026

[Bug] --enable-flashinfer-pure-allreduce is not available on 0.5.18 #37646

Open

5 tasks

jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026

[NVIDIA] Add flashinfer MNNVL backend for allreduce only (sgl-project…

192b90e

…#30700)

Klaud-Cold mentioned this pull request Sep 10, 2026

[Klaud Cold] Update dsr1-fp4-b300-sglang SGLang image to v0.5.19-cu130 / 将 dsr1-fp4-b300-sglang 的 SGLang 镜像更新至 v0.5.19-cu130 SemiAnalysisAI/InferenceX#2951

Closed

Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026

[NVIDIA] Add flashinfer MNNVL backend for allreduce only (sgl-project…

ae286b3

…#30700)

Klaud-Cold mentioned this pull request Sep 11, 2026

[Klaud Cold] Update dsr1-fp4-b200-sglang SGLang image to v0.5.19-cu130 / 将 dsr1-fp4-b200-sglang 的 SGLang 镜像更新至 v0.5.19-cu130 SemiAnalysisAI/InferenceX#2989

Open

Uh oh!

Conversation

wenscarl commented Jul 9, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Changes

flashinfer_comm_fusion.py — new flashinfer_allreduce()

parallel_state.py — dispatch hook + group tagging

server_args.py — new flag

model_runner.py — wiring

overrides.py — auto-enable for DeepSeek-V3/V4

test_flashinfer_pure_allreduce.py — unit tests

Performance and Accuracy

Uh oh!

gemini-code-assist Bot commented Jul 9, 2026

Uh oh!

nvpohanh commented Jul 15, 2026

Uh oh!

nvpohanh commented Jul 15, 2026

Uh oh!

shyeh25 commented Jul 15, 2026

Uh oh!

gemini-code-assist Bot commented Jul 16, 2026

Uh oh!

nvpohanh commented Jul 17, 2026

Uh oh!

nvpohanh commented Jul 17, 2026

Uh oh!

nvpohanh commented Jul 20, 2026

Uh oh!

ormandj commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nvpohanh commented Jul 28, 2026

Uh oh!

b8zhong commented Jul 29, 2026

Uh oh!

github-actions Bot commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

wenscarl commented Jul 9, 2026 •

edited by github-actions Bot

Loading

`flashinfer_comm_fusion.py` — new `flashinfer_allreduce()`

`parallel_state.py` — dispatch hook + group tagging

`server_args.py` — new flag

`model_runner.py` — wiring

`overrides.py` — auto-enable for DeepSeek-V3/V4

`test_flashinfer_pure_allreduce.py` — unit tests

ormandj commented Jul 24, 2026 •

edited

Loading

github-actions Bot commented Jul 29, 2026 •

edited

Loading

github-actions Bot commented Aug 12, 2026 •

edited

Loading