[MUSA][16/N] Add MUSA backend support for layers by popsiclexu · Pull Request #19509 · sgl-project/sglang

popsiclexu · 2026-02-27T12:40:04Z

Motivation

This PR is the 16th in a series of pull requests (tracked in #16565) to add full support for Moore Threads GPUs, leveraging MUSA (Meta-computing Unified System Architecture) to accelerate LLM inference.

dependencies： #17946

Modifications

Add forward_musa implementations for activation and quantization layers to support MUSA backend.
Update condition checks to include MUSA alongside CUDA/HIP/XPU for kernel imports and execution paths.
Modify deep_gemm integration to handle MUSA-specific imports and disable certain optimizations.
Add MUSA-specific implementations for sampling, layernorm, MoE, and FP8 linear operations.

Accuracy Tests

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

gemini-code-assist · 2026-02-27T12:40:33Z

Summary of Changes

Hello, 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 introduces comprehensive support for Moore Threads (MUSA) GPUs within the SGLang runtime (SRT). It enables efficient LLM inference on MUSA hardware by integrating MUSA-specific logic and implementations across critical components such as activation functions, normalization, DeepGEMM operations, Mixture-of-Experts layers, FP8 quantization, and sampling mechanisms.

Highlights

MUSA Backend Integration: Added forward_musa implementations for activation, layernorm, and unquantization layers to support the MUSA backend.
DeepGEMM Adaptation: Modified DeepGEMM integration to handle MUSA-specific imports, disable certain optimizations, and adjust scaling for FP8 operations.
MoE Layer Support: Extended Mixture-of-Experts (MoE) layers with MUSA-specific kernel imports, atomic add semantics, activation handling, and block size alignment.
FP8 Quantization & Linear Operations: Implemented MUSA-specific FP8 quantization kernels and linear operations, including a dedicated auto-backend dispatch for MUSA hardware.
Sampling Enhancements: Integrated MUSA backend support for sampling, including conditional kernel imports and adjustments for greedy sampling.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

python/sglang/srt/layers/activation.py
- Added forward_musa implementations for SwishGLU and GeluAndMul classes.
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
- Integrated MUSA-specific DeepGEMM imports and disabled DeepGEMM compilation for MUSA.
python/sglang/srt/layers/deep_gemm_wrapper/configurer.py
- Updated DeepGEMM configuration with MUSA capability checks and conditional imports.
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
- Adjusted DeepGEMM entrypoint for MUSA-specific imports and environment variable access.
python/sglang/srt/layers/layernorm.py
- Added MUSA support for RMS normalization, including conditional flashinfer import and a dedicated forward_musa method.
python/sglang/srt/layers/moe/ep_moe/kernels.py
- Modified MoE expert scattering kernel to support MUSA's atomic add semantics.
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
- Integrated MUSA support into fused MoE Triton kernels, including conditional imports, activation handling, and sum reduction.
python/sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py
- Enabled MUSA for MoE block size alignment.
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
- Adapted DeepGEMM MoE runner to handle MUSA-specific scaling and alignment.
python/sglang/srt/layers/moe/moe_runner/triton.py
- Extended Triton-based MoE runner with MUSA support for activation functions and intermediate cache management.
python/sglang/srt/layers/moe/topk.py
- Enabled MUSA for MoE top-k selection, including conditional kernel imports and grouped top-k logic.
python/sglang/srt/layers/quantization/fp8.py
- Updated FP8 quantization layer to include MUSA capability checks.
python/sglang/srt/layers/quantization/fp8_kernel.py
- Adapted FP8 quantization kernels for MUSA backend, including conditional imports.
python/sglang/srt/layers/quantization/fp8_utils.py
- Implemented MUSA-specific FP8 linear operations and integrated them into the auto-backend dispatch.
python/sglang/srt/layers/quantization/unquant.py
- Added MUSA forward method for unquantization, reusing CUDA implementation.
python/sglang/srt/layers/sampler.py
- Extended sampling layer with MUSA backend support, including conditional kernel imports and greedy sampling adjustments.

Activity

The pull request was created by popsiclexu to add MUSA backend support for various layers, as part of a larger effort to support Moore Threads GPUs.

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.

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 MUSA backend support for various layers in the sglang library, including activation, quantization, layernorm, MoE, and FP8 linear operations. The changes involve adding new forward_musa implementations, updating conditional checks for kernel imports and execution paths, and modifying deep_gemm integration to handle MUSA-specific requirements. The code has been reviewed, and several issues related to conditional checks, potential dead code, and code duplication have been identified and reported in the review comments.

gemini-code-assist · 2026-02-27T12:47:18Z

+        if not _is_musa:
+            intermediate_cache2 = torch.empty(
+                (total_tokens, N // 2),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )


The intermediate_cache2 tensor is only created when _is_musa is False. This means that the subsequent code that uses intermediate_cache2 will not be executed when _is_musa is True. Please ensure that the MUSA backend does not require this tensor or provide an alternative implementation.

gemini-code-assist · 2026-02-27T12:47:18Z

+        if not _is_musa:
+            intermediate_cache2 = torch.empty(
+                (M * topk_ids.shape[1], N // 2),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )


The intermediate_cache2 tensor is only created when _is_musa is False. This means that the subsequent code that uses intermediate_cache2 will not be executed when _is_musa is True. Please ensure that the MUSA backend does not require this tensor or provide an alternative implementation.

gemini-code-assist · 2026-02-27T12:47:18Z

+        elif not _is_musa:
            hidden_states_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
                hidden_states_scale
            )


The code path for not _is_musa uses deep_gemm_wrapper.get_mn_major_tma_aligned_tensor, while the else path (when _is_musa) does not have any code. This could lead to incorrect behavior or performance degradation on MUSA devices. Please ensure that the MUSA backend has a proper implementation for this case.

gemini-code-assist · 2026-02-27T12:47:18Z

+        if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 or not _is_musa:
            down_input_scale = tma_align_input_scale(down_input_scale)



The TMA alignment is skipped for MUSA. Please ensure that this is the correct behavior and that MUSA does not require TMA alignment for deep_gemm. If TMA alignment is required, consider providing a MUSA-specific implementation.

gemini-code-assist · 2026-02-27T12:47:19Z

+        if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 or not _is_musa:
            down_input_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
                down_input_scale


The TMA alignment is skipped for MUSA. Please ensure that this is the correct behavior and that MUSA does not require TMA alignment for deep_gemm. If TMA alignment is required, consider providing a MUSA-specific implementation.

gemini-code-assist · 2026-02-27T12:47:19Z

+        if _is_musa:
+            return 31


The minimum capability is set to 31 for MUSA. Please ensure that this is the correct minimum capability for MUSA to support FP8. If FP8 is not supported for MUSA, consider removing the is_musa check from this condition to avoid confusion.

gemini-code-assist · 2026-02-27T12:47:19Z

+    else:
+        from sgl_kernel import sgl_per_tensor_quant_fp8, sgl_per_token_quant_fp8


When _is_musa is true, sgl_per_tensor_quant_fp8 and sgl_per_token_quant_fp8 are imported from sgl_kernel. Please ensure that these kernels are optimized for the MUSA architecture. If not, consider providing MUSA-specific implementations.

gemini-code-assist · 2026-02-27T12:47:19Z

+    if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and not _is_musa:
        return deepgemm_w8a8_block_fp8_linear_with_fallback


The deep_gemm backend is disabled for MUSA. Please ensure that this is the correct behavior and that MUSA does not support deep_gemm. If deep_gemm is supported, consider providing a MUSA-specific implementation.

gemini-code-assist · 2026-02-27T12:47:19Z

+    assert input_scale is None
+    input_2d = input.view(-1, input.shape[-1])
+    output_shape = [*input.shape[:-1], weight.shape[0]]
+    if input_2d.shape[0] < 4:
+        output = musa_fused_gemv(
+            input_2d,
+            weight,
+            None,
+            weight_scale,
+        )
+    else:
+        q_input, x_scale = sglang_per_token_group_quant_fp8(
+            input_2d, block_size[1], column_major_scales=False
+        )
+        output = musa_mudnn_w8a8_scaled_mm(
+            q_input,
+            weight,
+            out_dtype=input_2d.dtype,
+            scale_a=x_scale,
+            scale_b=weight_scale,
+        )
+    if bias is not None:
+        output += bias
+    return output.to(dtype=input_2d.dtype).view(*output_shape)


Adding a new function musa_w8a8_block_fp8_linear for MUSA. This is good for providing MUSA-specific implementations. Please ensure that this function is optimized for the MUSA architecture and that it is thoroughly tested.

gemini-code-assist · 2026-02-27T12:47:19Z

+            if _is_musa and logits.shape[0] < 1:
+                batch_next_token_ids = torch.empty(
+                    [0], dtype=torch.int64, device=logits.device
+                )


This condition checks if _is_musa and logits.shape[0] < 1. If this condition is true, an empty tensor is created. Please ensure that this is the correct behavior and that it does not cause any issues in the subsequent code. Also, consider adding a comment explaining why this condition is necessary.

froststeam · 2026-04-07T07:08:49Z

        if sampling_info.is_all_greedy:
-            # Use torch.argmax if all requests use greedy sampling
-            batch_next_token_ids = torch.argmax(logits, -1)
+            # XXX (MUSA): protect against empty logits tensor


This can be removed to maintain consistency with the community.

froststeam · 2026-04-07T07:15:47Z

+_is_musa = is_musa()
+
+if _is_musa:
+    from sgl_kernel import musa_fused_gemv, musa_mudnn_w8a8_scaled_mm


Instead of _is_musa branch, using deep_gemm_wrapper.gemm_nt_f8f8bf16 could be a better option.

popsiclexu · 2026-04-14T08:19:10Z

move to #22774

popsiclexu requested review from AniZpZ, BBuf, Edwardf0t1, FlamingoPg, Fridge003, HaiShaw, Ying1123, ch-wan, ispobock and merrymercy as code owners February 27, 2026 12:40

github-actions Bot added the quant LLM Quantization label Feb 27, 2026

popsiclexu changed the title ~~feat(musa): add MUSA backend support for layers~~ [MUSA][11/N] Add MUSA backend support for layers Feb 27, 2026

gemini-code-assist Bot reviewed Feb 27, 2026

View reviewed changes

yeahdongcn added the mthreads label Feb 27, 2026

yeahdongcn changed the title ~~[MUSA][11/N] Add MUSA backend support for layers~~ [MUSA][16/N] Add MUSA backend support for layers Feb 28, 2026

yeahdongcn marked this pull request as draft February 28, 2026 00:41

yeahdongcn mentioned this pull request Feb 28, 2026

[Roadmap][Feature] Support Moore Threads (MUSA) GPU #16565

Open

2 tasks

popsiclexu force-pushed the feat/musa_layers branch 3 times, most recently from ebcfbe1 to 74e7019 Compare March 3, 2026 08:32

popsiclexu force-pushed the feat/musa_layers branch from 74e7019 to d8fd705 Compare March 24, 2026 08:43

froststeam reviewed Apr 7, 2026

View reviewed changes

popsiclexu and others added 2 commits April 13, 2026 16:35

feat: add MUSA support for custom op registration

cd54b0c

feat(musa): add MUSA backend support for layers

91993c2

popsiclexu force-pushed the feat/musa_layers branch from d8fd705 to cbbcb20 Compare April 13, 2026 09:30

github-actions Bot added the deepseek label Apr 13, 2026

feat: support deepseek on MUSA

e5fe399

popsiclexu force-pushed the feat/musa_layers branch from cbbcb20 to e5fe399 Compare April 13, 2026 14:13

popsiclexu closed this Apr 14, 2026

		if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 or not _is_musa:
		down_input_scale = tma_align_input_scale(down_input_scale)

		else:
		from sgl_kernel import sgl_per_tensor_quant_fp8, sgl_per_token_quant_fp8

		if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and not _is_musa:
		return deepgemm_w8a8_block_fp8_linear_with_fallback

Conversation

popsiclexu commented Feb 27, 2026 • edited by yeahdongcn Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Feb 27, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

froststeam Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

froststeam Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

popsiclexu commented Apr 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

popsiclexu commented Feb 27, 2026 •

edited by yeahdongcn

Loading