minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] hardware: support sage attention backend on MUSA (attn backend, 21/N) (#24752)

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-05-11 19:50:52 -07:00
committed by GitHub
parent 5495026a3b
commit 0a37d24e62
3 changed files with 22 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/README.md
+1 -1
View File
@@ -25,7 +25,7 @@ SGLang Diffusion supports AMD Instinct GPUs through ROCm. On AMD platforms, we u
### Moore Threads/MUSA Support
SGLang Diffusion supports Moore Threads GPUs (MTGPU) through the MUSA software stack. On MUSA platforms, we use the Torch SDPA backend for attention. See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions.
SGLang Diffusion supports Moore Threads GPUs (MTGPU) through the MUSA software stack. On MUSA platforms, we use FlashAttention (FA3) when available; also supports Sage Attention when installed; otherwise falls back to the Torch SDPA backend. See the [installation guide](https://github.com/sgl-project/sglang/tree/main/docs/diffusion/installation.md) for setup instructions.
### Apple MPS Support
python/sglang/multimodal_gen/runtime/platforms/musa.py
+18 -1
View File
@@ -160,6 +160,23 @@ class MusaPlatformBase(Platform):
if selected_backend == AttentionBackendEnum.TORCH_SDPA:
logger.info("Using Torch SDPA backend")
return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
elif selected_backend == AttentionBackendEnum.SAGE_ATTN:
try:
from sageattention import sageattn # noqa: F401
from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn import ( # noqa: F401
SageAttentionBackend,
)
logger.info("Using Sage Attention backend")
return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend"
except ImportError as e:
logger.info(e)
logger.info(
"Sage Attention backend is not installed (To install it, run `pip install sageattention>=0.1.0`). Falling back to Flash Attention."
)
target_backend = AttentionBackendEnum.FA
elif selected_backend in [
AttentionBackendEnum.FA,
]:
@@ -208,7 +225,7 @@ class MusaPlatformBase(Platform):
logger.info("Using Torch SDPA backend")
return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
logger.info("Using FlashAttention (FA3) backend on MUSA")
logger.info("Using FlashAttention (FA3) backend")
return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend"
@classmethod