[MUSA][9/N] Add FA3 attention backend support through MATE (MUSA AI Tensor Engine)

[froststeam]

Motivation

This PR fixes the Flash Attention backend support that was previously merged in PR #17985 but later reverted in PR #22002 due to a bug. The original commit 2373552 caused CI failures (see failed CI job).

Previously, the MUSA-adapted flash attention implementation had a bug in the _forward_extend_impl method. The code was missing a proper mechanism to select the correct kernel implementation based on the fa_impl_ver parameter, causing it to always use the default FA3 implementation regardless of the specified version.

Fix Applied

After rebasing to the latest main branch, the kernel selection logic has been refactored and moved to the FlashAttentionBackend.__init__ method. This ensures that the appropriate flash attention implementation is selected during initialization based on the fa_impl_ver parameter.

1. Moved kernel selection to __init__: The logic to select the correct flash attention kernel (including MUSA-specific implementations) is now handled in the FlashAttentionBackend.__init__ method, where two instance variables are initialized:

   • self.flash_attn_with_kvcache: For cached attention operations
   • self.flash_attn_varlen_func: For variable-length attention operations

2. Updated forward methods: Both _forward_extend_impl and _forward_decode_impl now use these instance variables instead of directly calling the default implementations, ensuring the correct kernel is used based on the initialized configuration.

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the https://docs.sglang.io/developer_guide/contribution_guide.html#format-code-with-pre-commit .
• Add unit tests according to the https://docs.sglang.io/developer_guide/contribution_guide.html#run-and-add-unit-tests .
• Update documentation according to https://docs.sglang.io/developer_guide/contribution_guide.html#write-documentations .
• Provide accuracy and speed benchmark results according to https://docs.sglang.io/developer_guide/contribution_guide.html#test-the-accuracy and https://docs.sglang.io/developer_guide/contribution_guide.html#benchmark-the-speed .
• Follow the SGLang code style https://docs.sglang.io/developer_guide/contribution_guide.html#code-style-guidance .

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the https://github.com/sgl-project/sglang/blob/main/.github/MAINTAINER.md#pull-request-merge-process .
2. Get approvals from https://github.com/sgl-project/sglang/blob/main/.github/CODEOWNERS and other reviewers.
3. Trigger CI tests with https://docs.sglang.io/developer_guide/contribution_guide.html#how-to-trigger-ci-tests or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

Related Links:

• Original merge PR: [MUSA][9/N] Add FA3 attention backend support through MATE (MUSA AI Tensor Engine) #17985
• Revert PR: Revert "[MUSA][9/N] Add FA3 attention backend support through MATE (MUSA AI Tensor Engine)" #22002 (commit efa7b2d)
• Failed CI job: https://github.com/sgl-project/sglang/actions/runs/23928333410/job/69789912493
• Fix commits: 0af5fe5

[gemini-code-assist]

Code Review

This pull request introduces support for the MUSA (Moore Threads GPU) hardware backend, specifically focusing on Flash Attention integration. It adds necessary dependencies, configuration parameters, and a new MUSA-specific attention module that wraps the mate library's flash attention functions. The implementation uses a thread-local context manager to automatically inject scheduler metadata into attention calls. Key changes include updates to the attention registry, the FlashAttentionBackend to handle MUSA-specific logic, and server argument adjustments for MUSA compatibility. Feedback highlights potential issues with global buffer safety in multi-GPU environments, metadata cache collisions due to non-unique keys, and the implications of ignoring cu_seqlens_k_new in the MUSA implementation.

[yeahdongcn]

I think it would be better to split this into two commits: one carrying over changes from the previous PR, and another fixing the regression in selecting FA kernels for different NVIDIA GPU architectures. This should make it easier for the SGLang core team to review.

[yeahdongcn]

The key updates to resolve the previous regression in FA3/FA4 kernel wiring on NVIDIA GPUs should be here, since upstream main now selects and stores the kernels during __init__.

[yeahdongcn]

I'll go ahead and approve this PR to trigger CI. Could @Fridge003 and @Kangyan-Zhou please take a final look? Thanks!

[yeahdongcn]

/tag-and-rerun-ci