[Feature] Enable PTPC FP8 for compressed tensors moe (aiter kernel)

[qichu-yun]

Motivation

Deepseek-R1 PTPC FP8 model for llm-compressor quantized model can be supported on MI308.

Accuracy Tests

Here is my script.

model=/data/models/Deepseek-R1-FP8-Dynamic
TP=8
EP=1

python3 -m sglang.launch_server \
    --model-path ${model} \
    --host localhost \
    --port 9000 \
    --tp-size ${TP} \
    --ep-size ${EP} \
    --trust-remote-code \
    --chunked-prefill-size 16384 \
    --mem-fraction-static 0.85 \
    --disable-radix-cache \
    --max-prefill-tokens 16384 \
    --cuda-graph-max-bs 128 \
    2>&1 | tee log.server.log &

Parallel strategy：TP8
acc:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.932|±  |0.0160|
|     |       |strict-match    |     5|exact_match|↑  |0.928|±  |0.0164|

Parallel strategy：TP8 + EP8
acc:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.967|±  |0.0112|
|     |       |strict-match    |     5|exact_match|↑  |0.967|±  |0.0112|

Benchmarking

Here is my script.

export SGLANG_USE_AITER=1
model=/mnt/raid0/models/deepseek-r1-FP8-Dynamic-from-BF16

input_tokens=3584
output_tokens=1024
max_concurrency=64
num_prompts=128

echo "bench model: ${model}"
echo "input tokens: ${input_tokens}"
echo "output tokens: ${output_tokens}"
echo "max concurrency: ${max_concurrency}"
echo "num prompts: ${num_prompts}"

python3 -m sglang.bench_serving \
    --host localhost \
    --port 9000 \
    --model ${model} \
    --dataset-name random \
    --random-input ${input_tokens} \
    --random-output ${output_tokens} \
    --random-range-ratio 1.0 \
    --max-concurrency ${max_concurrency} \
    --num-prompt ${num_prompts} \

Parallel strategy：TP8

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 64        
Successful requests:                     128       
Benchmark duration (s):                  152.51    
Total input tokens:                      458752    
Total input text tokens:                 458752    
Total input vision tokens:               0         
Total generated tokens:                  131072    
Total generated tokens (retokenized):    130687    
Request throughput (req/s):              0.84      
Input token throughput (tok/s):          3007.92   
Output token throughput (tok/s):         859.41    
Total token throughput (tok/s):          3867.33   
Concurrency:                             63.97     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   76217.67  
Median E2E Latency (ms):                 76215.98  
---------------Time to First Token----------------
Mean TTFT (ms):                          11564.82  
Median TTFT (ms):                        11734.60  
P99 TTFT (ms):                           20507.28  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          63.20     
Median TPOT (ms):                        63.04     
P99 TPOT (ms):                           73.21     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           63.20     
Median ITL (ms):                         54.44     
P95 ITL (ms):                            55.93     
P99 ITL (ms):                            56.77     
Max ITL (ms):                            19669.39  
==================================================

Parallel strategy：TP8 + EP8

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 64        
Successful requests:                     128       
Benchmark duration (s):                  155.88    
Total input tokens:                      458752    
Total input text tokens:                 458752    
Total input vision tokens:               0         
Total generated tokens:                  131072    
Total generated tokens (retokenized):    64131     
Request throughput (req/s):              0.82      
Input token throughput (tok/s):          2942.89   
Output token throughput (tok/s):         840.83    
Total token throughput (tok/s):          3783.72   
Concurrency:                             63.96     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   77898.83  
Median E2E Latency (ms):                 77897.51  
---------------Time to First Token----------------
Mean TTFT (ms):                          10522.66  
Median TTFT (ms):                        10886.52  
P99 TTFT (ms):                           18062.87  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          65.86     
Median TPOT (ms):                        65.48     
P99 TPOT (ms):                           75.02     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           65.86     
Median ITL (ms):                         58.54     
P95 ITL (ms):                            60.12     
P99 ITL (ms):                            61.09     
Max ITL (ms):                            17324.30  
==================================================

As you can see, TP8 has better performance than TP8+EP8, and we recommend running TP8 for the PTPC model.

[gemini-code-assist]

Summary of Changes

Hello @qichu-yun, 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 quantization capabilities by introducing support for Per-Token Per-Channel FP8 quantization within Mixture-of-Experts layers, powered by the AITer kernel. The primary goal is to enable the use of advanced quantized models like Deepseek-R1 PTPC FP8, improving efficiency and compatibility without compromising accuracy. The changes involve core modifications to how weights are processed and how MoE and linear operations are executed under specific quantization schemes.

Highlights

• PTPC FP8 Quantization for MoE: Enabled Per-Token Per-Channel (PTPC) FP8 quantization for Mixture-of-Experts (MoE) layers, specifically leveraging the AITer kernel. This allows support for Deepseek-R1 PTPC FP8 models.
• AITer Kernel Integration: Integrated AITer's fused MoE operations and FP8 linear computations, including conditional weight padding and shuffling based on environment variables, to optimize performance for quantized models.
• Deepseek Model Compatibility: Adjusted the DeepseekV2MoE initialization to bypass certain weight block size assertions when using AITer with 'compressed-tensors' quantization for PTPC models, ensuring compatibility.
• New FP8 Linear Function: Introduced a new apply_fp8_ptpc_linear function in fp8_utils.py that utilizes AITer's per-token quantization and GEMM functions for efficient FP8 linear operations.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 PR enables PTPC FP8 for compressed tensors MoE with the aiter kernel, which is a great feature for AMD GPUs. The changes look mostly correct, but I've found a few issues related to code duplication, typos, and incorrect logic in shape calculation and weight processing. Specifically, there's a critical bug in compressed_tensors_w8a8_fp8.py due to a misplaced refactoring that could lead to incorrect model behavior, and another one in fp8_utils.py regarding output shape calculation. I've provided suggestions to fix these issues and improve code maintainability.

[qichu-yun]

Could you please review this PR when you're free, thanks! ? @wuhuikx @zejunchen-zejun @HaiShaw

[HaiShaw]

Please address comment left

[HaiShaw]

/tag-and-rerun-ci 11/26

[HaiShaw]

L105-L113 is redundant to L86-L91 in one case, if so (not intended to do twice), then need to refactor code.

[qichu-yun]

I have organized the code uniformly under the “elif self.strategy == QuantizationStrategy.CHANNEL:” branch.

[HaiShaw]

AITER_MOE_PADDING_SIZE - undocumented ENV introduced, not have its description, neither tests in PR body.
Please document it, describe its purpose.

Note - see if this serves the purpose:
https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L58

[qichu-yun]

I have removed the env AITER_MOE_PADDING_SIZE. There is no need to use this env in this case.

[HaiShaw]

@qichu-yun why keep the torch.nn.functional.pad calls if nothing really to be padded?