Fix one bug in the grouped-gemm triton kernel

[ch-wan]

Motivation

When block quant is not in use, EPMoE returns incorrect responses.

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! Gemini here, providing a summary of this pull request to help everyone quickly understand the changes.

This PR, authored by ch-wan, aims to fix a bug in the grouped-gemm Triton kernel used within the EPMoE (Expert Parallel Mixture of Experts) layer. According to the description, the bug causes incorrect responses when block quantization is not being used. The core change involves correcting how the scaling factor (scale_a) is loaded within the kernel to ensure it uses the correct index based on the current block's position rather than just the expert ID.

Highlights

• Bug Fix: Addresses an issue in the EPMoE layer where incorrect responses are produced when block quantization is disabled.
• Triton Kernel Update: Modifies the grouped_gemm_triton_kernel to correctly load the scale_a scaling factor using the block's position (m_range_start + offs_am) instead of just the expert_id under specific conditions (use_fp8_w8a8 is true and group_k/group_n are not both positive).

Changelog

• python/sglang/srt/layers/moe/ep_moe/kernels.py
  • Corrected the indexing logic for loading the scale_a scaling factor in the grouped_gemm_triton_kernel function (line 624) when using FP8 weights/activations without specific grouping conditions. This change ensures the correct scale is applied to the accumulator based on the current block being processed.

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.

---

A kernel's tiny fix,
Scales loaded, no more tricks.
MoE runs true,
Responses new,
No bugs in the matrix.

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. ↩

[zhyncs]

When block quant is not in use

which use case?

[gemini-code-assist]

Code Review

This pull request aims to fix a bug in the grouped_gemm_triton_kernel related to incorrect responses when block quantization is not used in EPMoE. The change modifies how scale_a_value is loaded. While the intent is to fix a bug, the change introduces a potential critical issue if the scale_a tensor is not structured as per-token scales, and also highlights an inconsistency with how activation scales appear to be computed and used in the calling EPMoE layer.

Summary of Findings

• Potential Out-of-Bounds Read and Scale Consistency: The change to load scale_a_value using per-token indices (scale_a + m_range_start + offs_am[:, None]) in grouped_gemm_triton_kernel is correct if and only if the scale_a tensor is structured as per-token scales. If scale_a is a smaller per-expert tensor (as suggested by its current calculation in EPMoE.forward), this change could lead to critical out-of-bounds memory access. Furthermore, for the dequantization to be correct, the initial quantization of activations must also use corresponding per-token scales, which seems inconsistent with how pre_reorder_triton_kernel currently uses scale_a.

Merge Readiness

The pull request addresses an issue with incorrect responses, but the proposed change introduces a significant risk of out-of-bounds memory access if the scale_a tensor is not structured as per-token. There also appear to be inconsistencies with how activation scales are computed and used in related parts of the code.

Due to the critical nature of the potential OOB read and the identified inconsistencies, I recommend that these concerns be thoroughly investigated and addressed before merging. I am unable to approve the pull request in its current state; further review and validation by others are necessary.

[gemini-code-assist]

This change modifies scale_a_value to be loaded using per-token indexing (m_range_start + offs_am[:, None]). This implies that the scale_a tensor (which originates from self.w13_input_scale or self.w2_input_scale in layer.py) is expected to be a 1D tensor containing per-token scale factors.

However, there are a few concerns:

1. Potential Out-of-Bounds Access: If scale_a is actually a smaller, per-expert scale tensor (as suggested by its calculation in EPMoE.forward at layer.py#L262-L268, where it's derived from torch.max(hidden_states) and has shape (num_experts_per_partition,)), then indexing it with global token indices (m_range_start + offs_am) could lead to out-of-bounds memory access, which is a critical issue.
2. Consistency with Activation Quantization: For this per-token dequantization to be correct, the activation tensor a (input to this kernel) must have been quantized using corresponding per-token scales.
   • If a is pre-quantized to FP8 by pre_reorder_triton_kernel (as suggested by gateup_input dtype in layer.py#L252-L259), then pre_reorder_triton_kernel must use per-token scales. However, its current implementation (kernels.py#L163-L173) appears to use per-expert scales (tl.load(scale_a + expert_id_cur_rank)).
   • If a is not pre-quantized (e.g., it's bf16/fp16), then this kernel performs the quantization of a in its main loop (around kernels.py#L593-L603 in the full file). The same scale_a_value (now per-token) would be used for quantizing a and later for dequantizing the accumulator. This part would be consistent if scale_a is indeed per-token.

Could you clarify the structure of the scale_a tensor in this specific scenario (use_fp8_w8a8 and not (group_k > 0 and group_n > 0)) and ensure that it is compatible with per-token indexing? If scale_a is indeed intended to be per-token, the upstream code responsible for computing w13_input_scale/w2_input_scale and its usage in pre_reorder_triton_kernel might need adjustments to ensure consistency.

[Alcanderian]

It will cause illegal mem access in epmoefp8 test: https://github.com/sgl-project/sglang/actions/runs/15342843382/job/43197363952