[NVIDIA] Add flashinfer all-to-all MOE dispatcher

[trevor-m]

Draft PR since flashinfer-ai/flashinfer#2102 is not yet merged into flashinfer.

flashinfer-ai/flashinfer#2102 is now merged.

Motivation

This PR integrates the latest TRT-LLM moe all-to-all kernels into sglang (AKA nvlink one sided allotall or mnnvlthroughput alltoall):

NVLINK one-sided comm AllToAll strategy for throughput scenarios.

This implementation utilizes symmetric memory to enable peer-to-peer access between GPUs over NVLink.
The kernels only take the role as one side of the communication: the dispatch kernel puts the data
into peer ranks' symmetric memory from local buffer, while the combine kernel gets the data from peer
ranks' symmetric memory and reduces the data into local buffer. It is the most efficient implementation
by now, but requires symmetric memory size proportional to max_num_tokens * n_ranks, which may not
scale well for very large-scale parallelization.

Currently I have tested it with the flashinfer_cutlass moe runner backend and use fp4 quantize before communication. It also allows flashinfer_cutlass moe to write directly to the workspace buffer.

Remaining issues:

1. Need to improve max_num_tokens - should multiply the per_rank value by ep_size. Or calculate this automatically based on server args.
2. Currently flashinfer a2a doesn't support when ranks have 0 tokens, so we need to create a dummy token in those cases. Doing the tensor allocation prevents us from using streams to overlap with shared experts.

Modifications

• Add --moe-a2a-backend=flashinfer

Accuracy Tests

SGLANG_MOE_NVFP4_DISPATCH=1 python3 -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-v2 --trust-remote-code --quantization modelopt_fp4 --tp 4 --moe-runner-backend flashinfer_cutlass --ep-size 4 --dp 4 --enable-dp-attention --mem-fraction-static 0.85 --max-running-requests 2048 --stream-interval 5 --enable-dp-lm-head --attention-backend triton --cuda-graph-bs 1 2 4 8 16 32 64 128 256 512 --disable-radix-cache --moe-a2a-backend flashinfer
python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 2048 --random-input 1024 --random-output 1024 --random-range-ratio 1 --max-concurrency 2048
python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319 --port=30000
Accuracy: 0.955
Invalid: 0.000
Latency: 166.179 s
Output throughput: 874.407 token/s

Benchmarking and Profiling

Single node results are below. Working on multinode benchmarking

Single node 4xGB200 results and profiles

SGLANG_MOE_NVFP4_DISPATCH=1 python3 -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-v2 --trust-remote-code --quantization modelopt_fp4 --tp 4 --moe-runner-backend flashinfer_cutlass --ep-size 4 --dp 4 --enable-dp-attention --mem-fraction-static 0.85 --max-running-requests 2048 --stream-interval 5 --enable-dp-lm-head --cuda-graph-bs 1 2 4 8 16 32 64 128 256 512 --disable-radix-cache --moe-a2a-backend flashinfer --disable-shared-experts-fusion
python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 2048 --random-input 1024 --random-output 1024 --random-range-ratio 1 --max-concurrency 2048
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 2048      
Successful requests:                     2048      
Benchmark duration (s):                  105.78    
Total input tokens:                      2097152   
Total input text tokens:                 2097152   
Total input vision tokens:               0         
Total generated tokens:                  2097152   
Total generated tokens (retokenized):    2093741   
Request throughput (req/s):              19.36     
Input token throughput (tok/s):          19826.37  
Output token throughput (tok/s):         19826.37  
Peak output token throughput (tok/s):    41273.00  
Peak concurrent requests:                2048      
Total token throughput (tok/s):          39652.75  
Concurrency:                             2033.67   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   105036.00 
Median E2E Latency (ms):                 105036.65 
---------------Time to First Token----------------
Mean TTFT (ms):                          19192.10  
Median TTFT (ms):                        18945.35  
P99 TTFT (ms):                           35753.33  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          83.91     
Median TPOT (ms):                        84.31     
P99 TPOT (ms):                           96.69     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           83.91     
Median ITL (ms):                         64.86     
P95 ITL (ms):                            134.52    
P99 ITL (ms):                            187.70    
Max ITL (ms):                            6041.95   
==================================================

With SGLANG_MOE_NVFP4_DISPATCH=0 (BF16 dispatch)
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 2048      
Successful requests:                     2048      
Benchmark duration (s):                  113.61    
Total input tokens:                      2097152   
Total input text tokens:                 2097152   
Total input vision tokens:               0         
Total generated tokens:                  2097152   
Total generated tokens (retokenized):    2094097   
Request throughput (req/s):              18.03     
Input token throughput (tok/s):          18459.37  
Output token throughput (tok/s):         18459.37  
Peak output token throughput (tok/s):    39440.00  
Peak concurrent requests:                2048      
Total token throughput (tok/s):          36918.74  
Concurrency:                             2032.17   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   112730.76 
Median E2E Latency (ms):                 112745.83 
---------------Time to First Token----------------
Mean TTFT (ms):                          22471.12  
Median TTFT (ms):                        22192.53  
P99 TTFT (ms):                           42519.44  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          88.23     
Median TPOT (ms):                        88.62     
P99 TPOT (ms):                           104.26    
---------------Inter-Token Latency----------------
Mean ITL (ms):                           88.23     
Median ITL (ms):                         66.35     
P95 ITL (ms):                            139.80    
P99 ITL (ms):                            192.80    
Max ITL (ms):                            7388.84   
==================================================

With --moe-a2a-backend=none
============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 2048      
Successful requests:                     2048      
Benchmark duration (s):                  107.98    
Total input tokens:                      2097152   
Total input text tokens:                 2097152   
Total input vision tokens:               0         
Total generated tokens:                  2097152   
Total generated tokens (retokenized):    2093493   
Request throughput (req/s):              18.97     
Input token throughput (tok/s):          19420.79  
Output token throughput (tok/s):         19420.79  
Peak output token throughput (tok/s):    43370.00  
Peak concurrent requests:                2048      
Total token throughput (tok/s):          38841.58  
Concurrency:                             2036.25   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   107365.61 
Median E2E Latency (ms):                 107376.61 
---------------Time to First Token----------------
Mean TTFT (ms):                          19164.99  
Median TTFT (ms):                        18753.23  
P99 TTFT (ms):                           35803.22  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          86.22     
Median TPOT (ms):                        86.69     
P99 TPOT (ms):                           98.54     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           86.22     
Median ITL (ms):                         66.88     
P95 ITL (ms):                            149.64    
P99 ITL (ms):                            224.95    
Max ITL (ms):                            5954.79   
==================================================

Dispatch (bs=512) with --moe-a2a-backend=flashinfer

Dispatch (bs=512) with --moe-a2a-backend=none (FP4 allgather)

Combine (bs=512) with --moe-a2a-backend=flashinfer

Combine (bs-512) --moe-a2a-backend=none (reduce-scatter)

(bs=512) with --moe-a2a-backend=none (reduce-scatter)

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Warning

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

[ch-wan]

We also need to update this document: https://github.com/sgl-project/sglang/blob/94e1251131ca27260cb0e8938aeb7b4a4e630b19/docs/advanced_features/expert_parallelism.md#backends-for-all-to-all-communication

[fzyzcjy]

Hi, I am wondering whether this will be helpful for flashinfer_trtllm moe, and whether there will be support for it? Thanks!

[trevor-m]

Hi, I am wondering whether this will be helpful for flashinfer_trtllm moe, and whether there will be support for it? Thanks!

@fzyzcjy It looks like it should work based on NVIDIA/TensorRT-LLM@e4bf29b
I can give it a try

[ch-wan]

LGTM in general. Can we add a unittest for this new feature? Also, we need to update this document: https://github.com/sgl-project/sglang/blob/94e1251131ca27260cb0e8938aeb7b4a4e630b19/docs/advanced_features/expert_parallelism.md#backends-for-all-to-all-communication

[Fridge003]

/tag-and-rerun-ci

[fzyzcjy]

@trevor-m: I can give it a try

Hi, is there any updates about that? if I understand correctly trtllm moe should be used for decode, and thus this feature is most useful when combined with that

[Fridge003]

Can we move this test to test/srt/ep and register it at nightly test.
Can open a following PR for this

[trevor-m]

@trevor-m: I can give it a try

Hi, is there any updates about that? if I understand correctly trtllm moe should be used for decode, and thus this feature is most useful when combined with that

@fzyzcjy I tried this out. The problem is that currently flashinfer doesn't support doing trtllm moe without fused routing. For all-to-all, we need to do the routing separately first, then do the communication, then run moe. Here is my WIP branch for the sglang changes which would allow it otherwise: https://github.com/trevor-m/sglang/tree/trtllm-a2a-wip

Edit: I found that flashinfer has a separate api for this: "trtllm_fp4_block_scale_routed_moe". I will try it out.

I looked at our WideEP configs and it looks like flashinfer_cutedsl moe is actually what we use for decode (flashinfer_cutlass is used for prefill). Let me see if that can be enabled with this all-to-all.