[feat] Support tp mode for DeepSeek-R1-W4AFP8

[chenxijun1029]

Motivation

Support tp mode for DeepSeek w4a8 model, which has a better performace than ep mode.

Modifications

1. Add W4AFp8MoEMethod and associated create_weights, process_weights_after_loading function and apply function. In the apply function, we use the same cutlass_w4a8_moe kernel as ep moe uses.
2. Add some tile shape and cluster shape config for tp moe in cutlass_w4a8_moe kernel.
3. Add a router logic in w4afp8 quant config and method. When "enable_ep_moe" found in global_server_args_dict, we use ep mode, else tp.

Co-author: @yuhyao 827623970@qq.com

Benchmark

We run DeepSeek-R1-W4AFP8 on 8H20 with tp8, comparing to run DeepSeek-R1 on 8H20 with ep8.
Test configuration:

ISL1000, OSL1000

input/output len = 1000/1000, qps=128, max_concurrency=128, num_prompt=256.
The results are shown below:

TP

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    128.0     
Max request concurrency:                 128       
Successful requests:                     256       
Benchmark duration (s):                  159.00    
Total input tokens:                      256000    
Total generated tokens:                  256000    
Total generated tokens (retokenized):    254696    
Request throughput (req/s):              1.61      
Input token throughput (tok/s):          1610.09   
Output token throughput (tok/s):         1610.09   
Total token throughput (tok/s):          3220.18   
Concurrency:                             127.52    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   79201.47  
Median E2E Latency (ms):                 78956.21  
---------------Time to First Token----------------
Mean TTFT (ms):                          6547.98   
Median TTFT (ms):                        6612.11   
P99 TTFT (ms):                           11687.82  
---------------Inter-Token Latency----------------
Mean ITL (ms):                           72.73     
Median ITL (ms):                         68.05     
P95 ITL (ms):                            72.42     
P99 ITL (ms):                            73.05     
Max ITL (ms):                            11148.65  
==================================================

While EP:

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    128.0     
Max request concurrency:                 128       
Successful requests:                     256       
Benchmark duration (s):                  161.37    
Total input tokens:                      256000    
Total generated tokens:                  256000    
Total generated tokens (retokenized):    255343    
Request throughput (req/s):              1.59      
Input token throughput (tok/s):          1586.41   
Output token throughput (tok/s):         1586.41   
Total token throughput (tok/s):          3172.81   
Concurrency:                             127.53    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   80390.85  
Median E2E Latency (ms):                 80644.43  
---------------Time to First Token----------------
Mean TTFT (ms):                          8143.41   
Median TTFT (ms):                        8144.56   
P99 TTFT (ms):                           14833.05  
---------------Inter-Token Latency----------------
Mean ITL (ms):                           72.32     
Median ITL (ms):                         66.38     
P95 ITL (ms):                            71.43     
P99 ITL (ms):                            72.08     
Max ITL (ms):                            14113.65  
==================================================

ISL6000, OSL1000

input/output len = 6000/1000, qps=128, max_concurrency=128, num_prompt=256.
The results are shown below:

TP

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    128.0     
Max request concurrency:                 128       
Successful requests:                     256       
Benchmark duration (s):                  525.30    
Total input tokens:                      1536000   
Total generated tokens:                  256000    
Total generated tokens (retokenized):    254578    
Request throughput (req/s):              0.49      
Input token throughput (tok/s):          2924.03   
Output token throughput (tok/s):         487.34    
Total token throughput (tok/s):          3411.37   
Concurrency:                             109.09    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   223849.73 
Median E2E Latency (ms):                 251328.94 
---------------Time to First Token----------------
Mean TTFT (ms):                          118601.26 
Median TTFT (ms):                        149225.11 
P99 TTFT (ms):                           183862.13 
---------------Inter-Token Latency----------------
Mean ITL (ms):                           105.36    
Median ITL (ms):                         83.10     
P95 ITL (ms):                            85.79     
P99 ITL (ms):                            86.27     
Max ITL (ms):                            52034.42  
==================================================

While EP:

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    128.0     
Max request concurrency:                 128       
Successful requests:                     256       
Benchmark duration (s):                  592.58    
Total input tokens:                      1536000   
Total generated tokens:                  256000    
Total generated tokens (retokenized):    255781    
Request throughput (req/s):              0.43      
Input token throughput (tok/s):          2592.07   
Output token throughput (tok/s):         432.01    
Total token throughput (tok/s):          3024.08   
Concurrency:                             108.41    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   250942.46 
Median E2E Latency (ms):                 282080.60 
---------------Time to First Token----------------
Mean TTFT (ms):                          137965.32 
Median TTFT (ms):                        171540.14 
P99 TTFT (ms):                           217855.50 
---------------Inter-Token Latency----------------
Mean ITL (ms):                           113.09    
Median ITL (ms):                         81.36     
P95 ITL (ms):                            84.83     
P99 ITL (ms):                            85.55     
Max ITL (ms):                            71910.62  
==================================================

We conclude the results as the below sheet:

Scenario	Version	ITL (ms)	TTFT (ms)	Request throughput (req/s)
1K input，1K output，128 request rate	EP	72.32	8 143.41	1.59
	TP	72.73 (+0.6 %)	6 547.98 (−19.60 %)	1.61 (+1.2 %)
6K input，1K output，128 request rate	EP	113.09	137 965.32	0.43
	TP	105.36 (−6.83 %)	118 601.26 (−14.03 %)	0.49 (+14.0 %)

#Accuracy
mmlu: 86.9 （by @yuhyao ）
aime24: 80.0
gpqa: 71.2
math500: 95.6

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the [Running Unit Tests] - The unittest for cutlass_w4a8_moe kernel already exsited (https://docs.sglang.ai/references/contribution_guide.html#running-unit-tests-adding-to-ci).
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Summary of Changes

Hello @chenxijun1029, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the DeepSeek-R1-W4AFP8 model's performance by introducing support for Tensor Parallelism (TP) mode. It involves a comprehensive update to the quantization and MoE layer implementations, including new weight handling, dynamic method selection, and optimized low-level kernel configurations, ultimately aiming for more efficient model serving.

Highlights

• Tensor Parallelism for DeepSeek-R1-W4AFP8: Introduced a new W4AFp8TPMoEMethod to enable Tensor Parallelism (TP) for DeepSeek-R1-W4AFP8 models, demonstrating improved performance compared to Expert Parallelism (EP) mode.
• Dynamic MoE Quantization Method Selection: Implemented a routing mechanism that dynamically selects between Expert Parallelism (EP) and Tensor Parallelism (TP) quantization methods for MoE layers based on the enable_ep_moe global server argument.
• Optimized CUTLASS Kernel Configurations: Added new tile and cluster shape configurations within the cutlass_w4a8_moe kernel and extended its dispatch logic to provide optimized performance for the specific matrix dimensions encountered in Tensor Parallelism mode.
• Enhanced Weight Processing and Loading: Updated the weight creation and processing logic for TP MoE, including refactoring scale interleaving into a shared utility function and adding support for special naming rules for input scales in mixed-precision models.
• Shared Experts Fusion Disablement: Explicitly disabled shared experts fusion optimization for W4A8 TP MoE models, as their quantization methods differ between routed and shared experts.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds support for tensor parallelism (TP) mode for the DeepSeek-R1-W4AFP8 model, which shows performance improvements over the existing expert parallelism (EP) mode. The changes include adding a new W4AFp8TPMoEMethod for quantization, updating the MoE kernel configurations, and adding logic to switch between TP and EP modes. The implementation looks solid and the benchmark results are promising. I've added a couple of comments to improve code maintainability by reducing duplication in both the Python and CUDA C++ code. These changes should not affect performance but will make the code easier to read and maintain.

[gemini-code-assist]

There's significant code duplication in the new else if blocks for n == 512 && k == 7168 and n == 7168 && k == 256. The only difference between the if (m <= ...) blocks is the Cutlass3xW4A8GemmSelected type.

To improve maintainability, you could refactor this using a helper template function or a macro to reduce the repeated calls to cutlass_w4a8_group_gemm_caller.

For example:

template <typename Gemm>
void invoke_cutlass_caller(...) {
    cutlass_w4a8_group_gemm_caller<Gemm>(...);
}

// In dispatch_w4a8_moe_mm_sm90:
if (m <= 4) {
    using Gemm = typename JOIN_STRUCT_NAME(64, 32, 512, 2, 1, 1)::Cutlass3xW4A8Gemm;
    invoke_cutlass_caller<Gemm>(...);
} else if (m <= 16) {
    using Gemm = typename JOIN_STRUCT_NAME_CO(128, 16, 512, 2, 1, 1)::Cutlass3xW4A8Gemm;
    invoke_cutlass_caller<Gemm>(...);
}
// ...

While I understand this pattern is common for performance-critical CUDA code to help the compiler generate specialized code, this refactoring would make the code much cleaner without a performance penalty.

[AniZpZ]

Thanks for your great work on this! To help us evaluate the impact of this PR, could you please provide the performance results (like GSM8K, MMLU, and Hellaswag)

[yangsijia-serena]

just want to confirm if you've checked the contents of the act_scales.safetensors file. Are the input scales for w1 and w3 all consistent?

[yuhyao]

@chenxijun1029 Nice work! Just wondering if there will be any further updates?
Also, should the file sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py be updated as well?

[zhilingjiang]

Nice work!

[chenxijun1029]

Thanks for your great work on this! To help us evaluate the impact of this PR, could you please provide the performance results (like GSM8K, MMLU, and Hellaswag)

Hi, I have updated the accuracy in the pr, please check it.

[AniZpZ]

LGTM, do you have any more sugguestion? @yangsijia-serena

[yangsijia-serena]

LGTM, do you have any more sugguestion? @yangsijia-serena

LGTM!

[BBuf]

It broken compressed-tensor format moe models like neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8 , fixed in #10299 .

[Bruce-x-1997]

hello, thanks for the pr, but when I use w4afp8 in deepseek-r1(https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8) and test it in aime24, it could only get 10% score
so what's your test model, which reach 80% in aime24
@chenxijun1029

[chenxijun1029]

hello, thanks for the pr, but when I use w4afp8 in deepseek-r1(https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8) and test it in aime24, it could only get 10% score so what's your test model, which reach 80% in aime24 @chenxijun1029

I tested it with evalscope, and you should set "max_num_tokens" larger, i.e. 20000, for better performance.

[Bruce-x-1997]

hello, thanks for the pr, but when I use w4afp8 in deepseek-r1(https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8) and test it in aime24, it could only get 10% score so what's your test model, which reach 80% in aime24 @chenxijun1029

I tested it with evalscope, and you should set "max_num_tokens" larger, i.e. 20000, for better performance.

thanks, I have reached 80% in r1 & r1-0528(quant use modelopt)
but I found if I use the same method to v3.1 , v3.1 accuracy could not be accepted, in v3.1(10% in aime24), and if we use a question in aime24 to test v3.1(after quant, and batch is 1), its answer is wrong.
do you find similar problems?is there anything we can try? @chenxijun1029

[chenxijun1029]

hello, thanks for the pr, but when I use w4afp8 in deepseek-r1(https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8) and test it in aime24, it could only get 10% score so what's your test model, which reach 80% in aime24 @chenxijun1029

I tested it with evalscope, and you should set "max_num_tokens" larger, i.e. 20000, for better performance.

thanks, I have reached 80% in r1 & r1-0528(quant use modelopt) but I found if I use the same method to v3.1 , v3.1 accuracy could not be accepted, in v3.1(10% in aime24), and if we use a question in aime24 to test v3.1(after quant, and batch is 1), its answer is wrong. do you find similar problems?is there anything we can try? @chenxijun1029

Can we share more information via wechat?

[llc-kc]

@chenxijun1029 Hi, are the TP and EP results evaluated by W4AFP8 and FP8 respectively? Do you compare W4AFP8 with the original FP8 both in TP mode? Why does W4AFP8 significantly reduce the weight size but seem not to improve ITL? Thanks.

[llc-kc]

@chenxijun1029 do we need an extra launch parameter --quantization w4afp8?

[junliu-mde]

@chenxijun1029 Hi, are the TP and EP results evaluated by W4AFP8 and FP8 respectively? Do you compare W4AFP8 with the original FP8 both in TP mode? Why does W4AFP8 significantly reduce the weight size but seem not to improve ITL? Thanks.

I guess it's the dequant process has non-negligible overhead. There should be space for improvement in kernels.