[Perf] Use Triton instead of Torch for DeepGEMM Per Token Group Quant

[yewentao256]

Purpose

Follow up for #20087

Use Triton instead of Torch for Per Token Group Quant

Test

Accuracy

VLLM_USE_DEEP_GEMM=1 lm_eval   --model vllm   --model_args "pretrained=Qwen/Qwen3-30B-A3B-FP8,max_model_len=32768,enforce_eager=True"   --trust_remote_code   --tasks gsm8k   --num_fewshot 5   --batch_size auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8476|±  |0.0099|
|     |       |strict-match    |     5|exact_match|↑  |0.8855|±  |0.0088|
# No deepgemm
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.8347|±  |0.0102|
|     |       |strict-match    |     5|exact_match|↑  |0.8939|±  |0.0085|

Performance

Qwen3 One card

VLLM_USE_DEEP_GEMM=1 vllm bench throughput  --model Qwen/Qwen3-30B-A3B-FP8 --load-format dummy --input-len 1000 --output-len 100 --trust_remote_code --enforce-eager --enable-expert-parallel
Throughput: 23.85 requests/s, 26179.47 total tokens/s, 2384.89 output tokens/s
# Before update
Throughput: 20.70 requests/s, 22721.25 total tokens/s, 2069.85 output tokens/s

R1: Don't have enough cards to test currently, but guessing it could improve like 15%+

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Summary of Changes

Hello @yewentao256, 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 focuses on a performance optimization for per-token group quantization by transitioning its underlying implementation from a potentially Torch-based or DeepGEMM-dependent approach to a dedicated Triton kernel. This migration not only simplifies the quantization logic by removing a conditional fallback but also introduces specific optimizations tailored for Blackwell GPUs, leading to a measurable increase in model throughput.

Highlights

• Migration to Triton for Per Token Group Quantization: The core change involves replacing the existing per_token_group_cast_to_fp8 function with a new, optimized Triton-based per_token_group_quant_fp8 implementation. This change is applied across relevant MoE and quantization layers and tests.
• Blackwell GPU Optimization: The Triton kernel for per-token group quantization (fp8_utils.py) has been enhanced to include a specific scaling method (ue8m0) that is conditionally applied when running on Blackwell GPUs, leveraging hardware capabilities for improved performance.
• Streamlined Quantization Logic: The pull request simplifies the overall quantization pipeline by removing a conditional fallback mechanism (vllm/utils/deep_gemm.py and vllm/model_executor/layers/fused_moe/utils.py), ensuring that all per-token group quantization now consistently uses the optimized Triton implementation.
• Performance Improvement: The changes result in a significant throughput increase, as demonstrated by benchmarks showing an improvement from 20.70 to 23.85 requests/s for Qwen3 on a single card.

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 replaces a wrapper function for per-token group quantization with a more direct Triton-based implementation, improving code clarity and maintainability. The changes are well-contained and the performance gains are documented. I have a few suggestions to improve the code's readability and correctness. Overall, this is a solid contribution.

[mgoin]

Can you run a benchmark to make sure perf isn't affected? It is a small detail but the documentation says

Note that x and y are always evaluated regardless of the value of condition.
https://triton-lang.org/main/python-api/generated/triton.language.where.html

[yewentao256]

Thanks for the catch! It doesn't affect too much currently because the other branch is just a =scale_raw
But I think you are right, we should just avoid using tl.where
Here is the new result:
Throughput: 23.91 requests/s, 26241.13 total tokens/s, 2390.51 output tokens/s

[mgoin]

LGTM! Great to see the perf improvement