Motivation

Closes: #4434

We found that online quantization of w8a8 is not supported, #4434 directly load a fp16 model with --quantization w8a8_fp8 will load a model with all weights as zero hence the benchmark result is nearly 0 for GSM8K.

This PR added such support and showcased reasonable benchmark result across different scenarios.

Benchmark and Testing

GSM8K

Online Quantization

1. Online Quantize a non-quantized model with --quantization w8a8_fp8: 0.789

   ```bash python3 -m sglang.launch_server --model /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --trust-remote-code --quantization w8a8_fp8 python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319

   Accuracy: 0.789
   Invalid: 0.002
   Latency: 17.456 s
   Output throughput: 7721.450 token/s

   Accuracy: 0.789
   Invalid: 0.002
   Latency: 17.199 s
   Output throughput: 7876.126 token/s

   </details>
   
   

2. Online Quantize a non-quantized model with --quantization fp8: 0.789

   ```bash python3 -m sglang.launch_server --model /shared/public/models/Meta-Llama-3-8B-Instruct --quantization fp8 --trust-remote-code python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319

   Accuracy: 0.789
   Invalid: 0.001
   Latency: 17.623 s
   Output throughput: 7737.012 token/s

   Accuracy: 0.790
   Invalid: 0.002
   Latency: 17.356 s
   Output throughput: 7759.684 token/s

   </details>
   
   

3. No Online Quantization: directly serve a non-quantized model as baseline: 0.791

   ```bash python3 -m sglang.launch_server --model /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --trust-remote-code python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319

   Accuracy: 0.791
   Invalid: 0.001
   Latency: 19.865 s
   Output throughput: 6720.907 token/s

   Accuracy: 0.787
   Invalid: 0.001
   Latency: 19.489 s
   Output throughput: 6838.216 token/s

   </details>
   
   
   

Offline Quantization (As sanity check)

1. Directly serve a pre-quantized model from Neural Magic with Dynamic Quantization recipe: 0.784

   ```bash python3 -m sglang.launch_server --model /shared/public/models/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic --trust-remote-code

   Accuracy: 0.784
   Invalid: 0.002
   Latency: 17.558 s
   Output throughput: 7646.241 token/s

   </details>
   
   In this case, online serving will use vllm kernels
   
   
   

2. Serve a pre-quantized model from Neural Magic with Dynamic Quantization recipe + online quantization override --quantization w8a8_fp8: 0.795

   ```bash python3 -m sglang.launch_server --model /shared/public/models/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic --trust-remote-code --quantization w8a8_fp8

   Accuracy: 0.795
   Invalid: 0.001
   Latency: 17.199 s
   Output throughput: 7856.940 token/s

   </details>
   
   In this case, online serving will use sgl kernels which were optimized in recent versions.
   
   
   

MMLU

Online Quantization

1. Unquantized model with --quantization w8a8_fp8: 0.682

   ```bash python3 -m sglang.launch_server --model /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --trust-remote-code --quantization w8a8_fp8 Total latency: 41.558 Average accuracy: 0.682 ```

2. Unquantized model with --quantization fp8: 0.682

   ```bash python3 -m sglang.launch_server --model /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --trust-remote-code --quantization fp8

   Total latency: 42.350
   Average accuracy: 0.682

   </details>
   
   

3. Unquantized model without any online quantization (Baseline): 0.683

   ``` python3 -m sglang.launch_server --model /shared/public/models/meta-llama/Meta-Llama-3.1-8B-Instruct --trust-remote-code

   Total latency: 52.971
   Average accuracy: 0.683

   </details>
   
   
   

Offline Quantization (As sanity check)

1. Directly serve a pre-quantized model from Neural Magic with Dynamic Quantization recipe: 0.683

   ```bash python3 -m sglang.launch_server --model /shared/public/models/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic --trust-remote-code python /home/jobuser/sglang/benchmark/mmlu/bench_sglang.py --nsub 64

   Total latency: 41.670
   Average accuracy: 0.683

   </details>
   
   In this case, online serving will use vllm kernels
   
   
   

2. Serve a pre-quantized model from Neural Magic with Dynamic Quantization recipe + online quantization override --quantization w8a8_fp8: 0.682

   ```bash python3 -m sglang.launch_server --model /shared/public/models/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic --trust-remote-code --quantization w8a8_fp8 python /home/jobuser/sglang/benchmark/mmlu/bench_sglang.py --nsub 64

   Total latency: 40.486
   Average accuracy: 0.682

   </details>
   
   In this case, online serving will use sgl kernels which were optimized in recent versions.
   
   
   
   

Modifications

• sglang/python/sglang/srt/layers/quantization/w8a8_fp8.py: Support online quantization while make sure offline quantization works as expected
• Added test cases to ensure online quant only can perform as certain level, it's worth mention that the accuracy is higher than pure offline quantization. We believe our SGL Kernel for quant and matmul could deliver better precision.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• 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.