[Perf] Tune FlashInfer all-reduce selection on SM103

[GirasoleY]

Tune standalone FlashInfer all-reduce using conservative, topology-specific cutoffs on SM103.

Scope

• Keep the existing VLLM_ALLREDUCE_USE_FLASHINFER gate and its default unchanged.
• When enabled, try FlashInfer before NCCL symmetric memory for inputs that satisfy FlashInfer's eligibility checks.
• Store standalone MNNVL tuning as (compute capability, world size, node count) in the shared all-reduce policy table.
• Use strict per-rank payload limits of 80 MiB for single-node TP4, 64 MiB for two-node TP8, and 8 MiB for four-node TP16.
• Keep TP2 on its existing 64 MiB default; the sweep validates it but does not require an override.
• Keep TP12 on the existing fallback. FlashInfer MNNVL rejects world size 12 at runtime; its supported sizes are {2, 4, 8, 16, 32, 64}.
• Promote the shared workspace to a tuned capacity only when standalone FlashInfer is enabled on a matching topology, keeping fusion-first CUDA graph pointers stable.
• Preserve existing thresholds and workspace sizing on every other topology and backend.

The separate fused all-reduce + RMSNorm threshold is intentionally excluded from this PR.

Performance

Hardware/software: GB300 GPUs in one NVL72 domain, PyTorch 2.13 / CUDA 13.0, FlashInfer 0.6.15, and NCCL 2.30.7. Values are median slowest-rank CUDA times and include the steady-state glue required by each backend. The fallback is consumer-visible NCCL symmetric memory for TP4/TP8/TP16 and ordinary out-of-place PyNCCL for TP2, where the symmetric-memory selector is disabled.

Topology	Strict cutoff	Edge payload	FlashInfer AUTO	Fallback	Fallback / FI - 1
TP2, 1 node	64 MiB (existing)	63.998 MiB	146.95 us	149.75 us	1.9%
TP4, 1 node	80 MiB	79.994 MiB	209.26 us	237.86 us	13.7%
TP8, 2 nodes	64 MiB	63.998 MiB	180.87 us	183.79 us	1.6%
TP12, 3 nodes	unsupported	4.006 MiB	rejected	111.04 us	n/a
TP16, 4 nodes	8 MiB	7.998 MiB	36.92 us	37.34 us	1.2%

The dense boundary sweeps used production-sized persistent workspaces. TP2 reaches parity and then loses with an 80 MiB workspace, so it retains 64 MiB. TP4 wins through 80 MiB, while a 96 MiB workspace loses at the 80 MiB point, so the cutoff stops at 80 MiB. TP16 wins through 8 MiB but loses clearly by 10 MiB, so it uses the strict 8 MiB cutoff.

K3 TP8 + DSpark decode

Three seed-matched measured cohorts per arm and concurrency, after two warmup waves. The setup used TP8, DCP1, DSpark7, FP8 KV cache, exact 1,024-token prompt KV state from DecodeBenchConnector, and 4,096 generated tokens per request.

Concurrency	Target AR activation	2 MiB TPOT	64 MiB TPOT	TPOT gain	Throughput gain
16	1.750 MiB	14.840 ms	14.833 ms	0.05%	0.01%
32	3.500 MiB	17.833 ms	17.167 ms	3.74%	3.66%

C16 is a control because its target activation is already below the old 2 MiB limit. C32 crosses that limit; after subtracting the C16 control shift, its threshold-associated TPOT benefit is 3.69 percentage points with a paired 95% interval of [3.49, 3.89] points.

All 288 measured requests completed with no failures. Acceptance lengths were identical between policies for every paired seed. FlashInfer was enabled and NCCL symmetric-memory all-reduce was disabled in both serving arms to isolate the threshold change; the full-glue sweep above measures the FlashInfer/fallback crossover directly.

Tests

• pytest -q tests/distributed/test_comm_ops.py -k flashinfer: 13 passed, 12 deselected.
• File-scoped pre-commit run on all four changed files: all applicable hooks passed.
• TP2/TP4/TP16 CUDA-graph sweeps checked numerical correctness and input preservation before and after replay.
• TP12 eager smoke test confirmed one-shot and two-shot both reject world size 12 while the fallback completes correctly.

Duplicate work and assistance

An open-PR search found #44080, which adjusts the fused all-reduce + RMSNorm cap on SM90/Hopper. This PR changes standalone MNNVL all-reduce selection and workspace sizing on SM103/GB300; it does not modify the fusion threshold. No duplicate open PR was found.

AI assistance was used to port, benchmark, test, and review this change. The submitter should review every changed line and the performance evidence before merge.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.