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

python/pyproject_other.toml

@@ -114,8 +114,11 @@ srt_musa = [
     "sglang[runtime_common]",
     "torch",
     "torch_musa",
-    "torchada>=0.1.45",
+    "torchada>=0.1.48",
     "mthreads-ml-py",
+    "mate",
+    "mate-deep_gemm",
+    "mate-flash-attention",
     "numpy<2.0",
 ]
 

python/sglang/srt/configs/model_config.py

@@ -575,6 +575,12 @@ def _derive_model_shapes(self):
         self.num_key_value_heads = getattr(
             self.hf_text_config, "num_key_value_heads", None
         )
+        self.first_k_dense_replace = getattr(
+            self.hf_text_config, "first_k_dense_replace", None
+        )
+        self.full_attention_interval = getattr(
+            self.hf_text_config, "full_attention_interval", None
+        )
 
         # for Dbrx and MPT models
         if self.hf_config.model_type in ["dbrx", "mpt"]:

python/sglang/srt/environ.py

@@ -331,6 +331,9 @@ class Envs:
     SGLANG_USE_AG_AFTER_QLORA = EnvBool(False)
     SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1)
 
+    # MTHREADS & MUSA
+    SGLANG_MUSA_FA3_FORCE_UPDATE_METADATA = EnvBool(False)
+
     # Quantization
     SGLANG_INT4_WEIGHT = EnvBool(False)
     SGLANG_CPU_QUANTIZATION = EnvBool(False)

python/sglang/srt/hardware_backend/musa/__init__.py

@@ -0,0 +1 @@
+# MUSA (Moore Threads GPU) hardware backend

python/sglang/srt/hardware_backend/musa/attention/__init__.py

@@ -0,0 +1,3 @@
+from .flashattention_backend import MusaFlashAttentionBackend
+
+__all__ = ["MusaFlashAttentionBackend"]