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

[ormandj]

Motivation

On SM12X (SM120/SM121) with TP>1, _resolve_backend rejects the FlashInfer all-reduce backend, so decode uses NCCL ring all-reduce. We reproduced this on SM120 (RTX PRO 6000 Blackwell). FlashInfer PR flashinfer-ai/flashinfer#3903 provides the SM12x TRT-LLM all-reduce kernel; with it installed, two runtime changes enable the path:

1. Extend the backend selector and the auto-enable override to SM120/SM121 (auto resolves to trtllm on single-node SM120/SM121, matching SM90).
2. Remove the multicast-granularity query from the TRT-LLM workspace preflight. The TRT-LLM workspace is non-multicast; the query is unnecessary and rejects valid non-multicast configurations.

Without FlashInfer #3903 installed, behavior is unchanged (fails closed exactly as today). No behavior change on SM90/SM100, covered by tests. FlashInfer flashinfer-ai/flashinfer#3930 was also present in both benchmark cells but is not required on current main — #30580 (lazy TileLang loading) already avoids the stub-library issue it addresses; it only matters for configurations that load TileLang eagerly before IPC initialization. Stacks on #30700, which provides the DeepSeek-V4 routing into this backend; happy to fold these commits into #30700 if maintainers prefer.

Modifications

• python/sglang/srt/layers/flashinfer_comm_fusion.py: accept SM12X (SM120/SM121) in _resolve_backend; drop the multicast-only granularity probe from _flashinfer_trtllm_workspace_allocation_sizes.
• python/sglang/srt/arg_groups/overrides.py: include SM12X in _flashinfer_allreduce_fusion_auto_enable.
• python/sglang/srt/server_args.py and docs_new/docs/advanced_features/server_arguments.mdx: document SM120/SM121 availability and backend selection.
• Tests (test/registered/unit/): selector resolution on SM12X, auto-enable on SM12X, and a regression test that the TRT-LLM workspace preflight does not require multicast support.

Accuracy Tests

GSM8K few-shot, 200 questions, identical config, NCCL cell vs this-PR cell: 97.5% vs 97.0% (0 invalid in both) — within the documented 1–5% batching variance. Greedy completions are identical through full code-generation spans across cells, diverging only at a free-form continuation boundary; expected, since one-shot Lamport changes floating-point reduction order vs ring. test_flashinfer_comm_fusion.py: 9 passed (14 subtests) in the benchmark container.

Speed Tests and Profiling

2x RTX PRO 6000 Blackwell (SM120, 96 GB, PCIe4 x16), TP2. DeepSeek-V4-Flash, FP8 KV cache, FlashInfer MXFP4 MoE, MTP2 (EAGLE steps=2, topk=1, draft-tokens=3), CUDA graphs, temperature 0. Both cells: SGLang main 3d91a569ce + PR #30700 (22e2f8b30d), FlashInfer 0.6.14 + PR #3903 (1eb77ab2) + PR #3930 (e855cc25). The only difference between cells is this PR's two runtime changes.

Cell	All-reduce path	C1 samples (tok/s)	median	mean
main + #30700	NCCL ring (#30700 fails closed on SM120)	125.90, 127.66, 130.19, 126.25, 118.75	126.25	125.75
main + #30700 + this PR	FlashInfer TRT-LLM one-shot Lamport	129.55, 131.81, 132.85, 140.27, 137.05	132.85	134.31

+5.2% on medians, +6.8% on means. Five consecutive 30 s runs per cell on a settled server after a warmup run. (An earlier revision of this body quoted single-run figures of 129.53 and 134.14; run-to-run drift on this setup reaches 6.9%, so those are superseded by the samples above.)

Torch profiler, median target-verify decode graph, TP0: 87 all-reduce kernels per graph in both cells; summed all-reduce time 0.772 ms (NCCL) → 0.394 ms (Lamport); graph GPU span 16.408 → 15.699 ms. Whole capture: 1,113 NCCL all-reduce kernels (176.5 ms) replaced 1:1 by 1,113 Lamport kernels (27.4 ms). Zero request errors and zero restarts across all runs. On a separate tuned stack (main 1e10ec93 + Torch 2.13 + PRs #29927/#30949/#32090/#32194/#32220 + #30700), per-graph all-reduce time was 0.695 → 0.367 ms; the end-to-end figures from that stack were single runs and are omitted here.

Checklist

• Format your code according to pre-commit (all hooks pass on changed files).
• Add unit tests (registered unit tree, see Modifications).
• Provide accuracy and speed benchmark results (above).
• Follow the SGLang code style guidance.

This contribution was developed with AI assistance. The author has reviewed the changes and their rationale, and all benchmark, profiling, and accuracy measurements were collected on the stated hardware.

---

CI States

Latest PR Test (Base): ❌ Run #30923777927
Latest PR Test (Extra): ❌ Run #30923776051

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[ormandj]

Currently not supported

This PR was intended to address this, is there no desire to support it? I apologize the PR body was a WIP, I should have done this as a draft first, but it is correctly updated now.

[mmangkad]

The original PR description was completely unrelated to the changes, so it looked suspicious, and I wasn't aware of the FI PR. This looks fine now, but it should wait until that PR lands and is in a stable FI release.

[ormandj]

The original PR description was completely unrelated to the changes, so it looked suspicious, and I wasn't aware of the FI PR. This looks fine now, but it should wait until that PR lands and is in a stable FI release.

I agree, I will keep an eye on things and follow the other PR and keep this updated. My apologies again, and thank you.

[ormandj]

Rebased onto current main; the conflict was the docs_new/ → docs/ rename. One commit added beyond the rebase: the workspace preflight retains multicast-granularity sizing for the mnnvl backend, while trtllm continues to skip the multicast query (which fails on SM120). Unit tests pass and the trtllm workspace initializes on both ranks on an SM120 TP2 deployment.

[ormandj]

Closing this proposal. Auto-selecting the TRT-LLM all-reduce backend for single-node SM120/SM121 is not optimal across all PCIe-only tensor-parallel configurations: TP4 testing showed a substantial regression that was removed by disabling FlashInfer all-reduce fusion. Any future change should be explicit opt-in or topology/TP-gated and validated separately at TP2, TP4, and TP8.