Support token-level quantization for EP MoE

[ch-wan]

Motivation

Modifications

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.

[gemini-code-assist]

Hello @ch-wan, 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!

Summary of Changes

Hello team, gemini-code-assist here to provide a summary of this pull request. This PR introduces support for token-level quantization within the Expert Parallel (EP) implementation of Mixture-of-Experts (MoE) layers. The primary goal is to enable more granular quantization, specifically for dynamic FP8 activation quantization, by calculating and applying scales on a per-token basis rather than per-expert or globally. This involves modifications to the underlying Triton kernels used for pre-reordering and grouped GEMM operations, as well as updates to the EPMoELayer and GroupedGemmRunner classes to handle and pass these per-token scales.

Highlights

• Token-Level Quantization: Adds the capability to perform quantization, particularly dynamic FP8 activation quantization, at the individual token level for EP MoE layers.
• Triton Kernel Modifications: Updates the pre_reorder_triton_kernel and grouped_gemm_triton_kernel to accept and utilize per-token scales when the use_per_token_if_dynamic flag is enabled.
• Per-Token Scale Calculation: Modifies the EPMoELayer to calculate per-token maximum values for dynamic activation quantization and derive per-token scales (w13_input_scale, w2_input_scale).
• Conditional Logic: Introduces a use_per_token_if_dynamic boolean flag throughout the relevant classes and kernels to conditionally enable or disable the per-token quantization behavior.

Changelog

Click here to see the changelog

• python/sglang/srt/layers/moe/ep_moe/kernels.py
  • Added use_per_token_if_dynamic constant parameter to pre_reorder_triton_kernel (Diff 1).
  • Modified scale loading in pre_reorder_triton_kernel to load per-token scale if enabled (Diff 2).
  • Added use_per_token_if_dynamic constant parameter to grouped_gemm_triton_kernel (Diff 3).
  • Modified scale_a loading in grouped_gemm_triton_kernel to load per-token scale if enabled for FP8 (Diff 4).
  • Added use_per_token_if_dynamic parameter to grouped_gemm_triton function (Diff 5).
  • Added assertion in grouped_gemm_triton to check scale shape when per-token is enabled (Diff 6).
  • Passed use_per_token_if_dynamic to kernel call in grouped_gemm_triton (Diff 7).
• python/sglang/srt/layers/moe/ep_moe/layer.py
  • Imported sglang_per_token_quant_fp8 (Diff 1).
  • Added use_per_token_if_dynamic parameter to GroupedGemmRunner constructor (Diff 2).
  • Passed use_per_token_if_dynamic to grouped_gemm_triton call in GroupedGemmRunner.forward (Diff 3).
  • Added use_per_token_if_dynamic parameter to EPMoELayer constructor (Diff 4, Diff 5).
  • Passed use_per_token_if_dynamic to GroupedGemmRunner constructor call in EPMoELayer.forward (Diff 6).
  • Modified w13_input_scale calculation for dynamic quantization to be per-token if enabled (Diff 7).
  • Added logic to reorder and assign per-token w13_input_scale after pre-reorder if enabled (Diff 8).
  • Modified down_input initialization and w2_input_scale handling for dynamic per-token quantization (Diff 9).
  • Added call to sglang_per_token_quant_fp8 to quantize down_input and get w2_input_scale if dynamic per-token is enabled (Diff 10).

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.

---

Quantize by the token,
Speed up the MoE,
Scales for each input,
Faster inference, you see.

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 introduces support for token-level quantization in the EP MoE layer, specifically for dynamic activation scaling in FP8. The changes involve adding a new flag use_per_token_if_dynamic to control this behavior in both the Python layer logic and the underlying Triton kernels. The implementation correctly branches based on this flag to apply per-token scaling where appropriate.

The core logic for calculating and applying per-token scales seems correctly implemented within the diffs. However, the introduction of this new mode adds complexity to the existing code paths, particularly in the EPMoE.forward method and the Triton kernels. Thorough testing and clear documentation are crucial for maintaining the quality and understanding of these changes.

Overall, this is a valuable feature addition, but there are several areas related to testing, documentation, and code maintainability that should be addressed before merging.

Summary of Findings

• Lack of Unit Tests: The pull request description indicates that unit tests should be added, but the corresponding checklist item is unchecked. New features, especially those involving performance-sensitive kernels and complex logic like per-token quantization, require comprehensive unit tests to ensure correctness and prevent regressions. This is a high-priority item.
• Increased Code Complexity: The introduction of the use_per_token_if_dynamic flag and the associated conditional logic in EPMoE.forward and the Triton kernels significantly increases the complexity of the code. This impacts maintainability and readability. While some complexity is inherent to adding a new quantization mode, consider if any parts can be refactored or simplified, perhaps by unifying data structures or kernel interfaces where possible.
• Potential Kernel Performance Implications: Conditional memory loads and branching within performance-critical loops in the Triton kernels (pre_reorder_triton_kernel, grouped_gemm_triton_kernel) might have performance implications. While Triton is capable of handling branches, minimizing them in hot loops is generally beneficial. Further profiling and potential kernel optimization might be needed.
• Lack of Documentation: The pull request description indicates that documentation should be updated, but the corresponding checklist item is unchecked. The new use_per_token_if_dynamic flag, its purpose, and how it affects the quantization behavior should be clearly documented, both in code comments/docstrings and potentially in user-facing documentation.

Merge Readiness

This pull request introduces a valuable new feature for token-level quantization. However, based on the review, there are high-priority issues related to the lack of unit tests and medium-priority issues concerning increased code complexity and lack of documentation. I recommend addressing these findings before merging. Please ensure comprehensive tests are added to cover the new logic and kernels, and update the documentation to reflect the changes. I am unable to approve this pull request; please have other reviewers assess the changes and approve before merging.

[gemini-code-assist]

This conditional load is placed inside the main loop iterating over topk. While Triton can handle branching, placing a conditional memory load inside a performance-critical loop like this might potentially impact performance or increase register pressure compared to loading outside the loop if possible. Could this load be moved outside the loop, perhaps by pre-calculating the scale for all tokens before the kernel launch if use_per_token_if_dynamic is true?

[gemini-code-assist]

Similar to the previous comment, this conditional check and load are inside the topk loop. While necessary for the logic, consider the potential performance implications of branching and memory access patterns within this loop.

[gemini-code-assist]

This conditional load of scale_a_value based on use_per_token_if_dynamic adds branching within the post-GEMM scaling logic. While this block is outside the main accumulation loop, it still adds complexity. Is there a way to unify the scale loading or application logic slightly, perhaps by ensuring scale_a always has the expected shape/layout for the chosen mode before calling the kernel?

[gemini-code-assist]

This assertion is a good check, but the logic for handling scale_a strides in the kernel call below (lines 734-735) uses ternary operators based on scale_a.ndim. This suggests the scale_a tensor might have different shapes/dimensions depending on whether per-token or per-expert scaling is used. This adds complexity to the kernel interface and the calling code. Could the scale_a tensor always be reshaped or prepared into a consistent format before calling the kernel, simplifying the kernel's stride logic?

[gemini-code-assist]

The logic for calculating w13_input_scale now has two distinct paths based on self.use_per_token_if_dynamic. This adds conditional complexity to the forward pass. While necessary for the feature, ensuring this logic is clearly documented and tested is important.

[gemini-code-assist]

This block adds logic to create and populate the w13_input_scale tensor specifically for the per-token dynamic case after the pre_reorder_triton_kernel. This is necessary to align the per-token scales with the reordered tokens using src2dst. This adds significant complexity to the forward pass state management. Could this scale preparation potentially be integrated into the pre_reorder_triton_kernel itself, or perhaps handled earlier?

[gemini-code-assist]

The creation and dtype of down_input are now conditional based on the quantization scheme and the new flag. This adds another layer of branching and complexity to the forward pass setup.

[gemini-code-assist]

This block adds conditional logic to perform per-token quantization on down_input and set self.w2_input_scale after the activation function. This is a core part of the new feature but further increases the complexity of the forward pass control flow.