feat: support sparse attention (#81)

Support sparse attention for vllm-plugin-fl.

Author: Zhangmj0621

vllm_fl/dispatch/backends/flaggems/flaggems.py

@@ -118,12 +118,13 @@ def rotary_embedding(
             inplace=inplace,
         )
 
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for FlagGems.
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use Deepseek Sparse Attention (DSA)
 
         Returns:
             Fully qualified class path string
@@ -140,4 +141,7 @@ def attention_backend(self, use_mla: bool = False) -> str:
         if use_mla:
             raise NotImplementedError("NOT support mla now!")
 
+        if use_sparse:
+            raise ValueError("use_sparse=True requires use_mla=True.")  
+
         return AttentionBackendEnum.TRITON_ATTN.get_path()

vllm_fl/dispatch/backends/reference/reference.py

@@ -120,7 +120,7 @@ def rotary_embedding(
             inplace=inplace,
         )
 
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for reference (vLLM native).
 
@@ -129,6 +129,7 @@ def attention_backend(self, use_mla: bool = False) -> str:
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use Deepseek Sparse Attention (DSA)
 
         Returns:
             Fully qualified class path string (vLLM native backend)
@@ -138,5 +139,7 @@ def attention_backend(self, use_mla: bool = False) -> str:
 
         if use_mla:
             # vLLM native MLA backend
+            if use_sparse:
+                return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
             return AttentionBackendEnum.FLASHMLA.get_path()
         return AttentionBackendEnum.FLASH_ATTN.get_path()

vllm_fl/dispatch/backends/vendor/ascend/ascend.py

@@ -127,7 +127,7 @@ def rotary_embedding(
             inplace=inplace,
         )
 
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for Ascend NPU.
 
@@ -140,10 +140,13 @@ def attention_backend(self, use_mla: bool = False) -> str:
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use Deepseek Sparse Attention (DSA)
 
         Returns:
             Fully qualified class path string
         """
         if use_mla:
+            if use_sparse:
+                raise NotImplementedError("MLA with sparse attention is not implemented for Ascend yet.")
             return "vllm_fl.dispatch.backends.vendor.ascend.impl.attention.AscendMLABackend"
         return "vllm_fl.dispatch.backends.vendor.ascend.impl.attention.AscendAttentionBackend"

vllm_fl/dispatch/backends/vendor/cuda/cuda.py

@@ -137,16 +137,18 @@ def rotary_embedding(
             inplace=inplace,
         )
 
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for CUDA.
 
         Supports:
         - FLASH_ATTN (default)
         - TRITON_ATTN (when use_flaggems_op("triton_attn") is True)
+        - FLASHMLA_SPARSE (when use_mla and use_sparse are both True)
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use Deepseek Sparse Attention (DSA)  
 
         Returns:
             Fully qualified class path string
@@ -155,6 +157,8 @@ def attention_backend(self, use_mla: bool = False) -> str:
         from vllm_fl.utils import use_flaggems_op
 
         if use_mla:
+            if use_sparse:
+                return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
             return AttentionBackendEnum.FLASHMLA.get_path()
 
         # Default to FLASH_ATTN

vllm_fl/dispatch/backends/vendor/iluvatar/iluvatar.py

@@ -143,19 +143,22 @@ def rotary_embedding(
             inplace=inplace,
         )
 
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for Iluvatar.
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use Deepseek Sparse Attention (DSA)
 
         Returns:
             Fully qualified class path string
         """
         from vllm.attention.backends.registry import AttentionBackendEnum
 
         if use_mla:
+            if use_sparse:
+                return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
             return AttentionBackendEnum.FLASHMLA.get_path()
 
         return AttentionBackendEnum.FLASH_ATTN.get_path()

vllm_fl/dispatch/backends/vendor/metax/metax.py

@@ -131,19 +131,23 @@ def rotary_embedding(
             inplace=inplace,
         )
 
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for METAX.
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use Deepseek Sparse Attention (DSA)
 
         Returns:
             Fully qualified class path string
         """
         from vllm.attention.backends.registry import AttentionBackendEnum
 
         if use_mla:
+            if use_sparse:
+                # TODO: Implement METAX MLA Sparse backend
+                return AttentionBackendEnum.FLASHMLA_SPARSE.get_path()
             # TODO: Implement METAX MLA backend
             return AttentionBackendEnum.FLASHMLA.get_path()
 

vllm_fl/dispatch/ops.py

@@ -129,7 +129,7 @@ def rotary_embedding(
     # ==================== Attention Backend ====================
 
     @abstractmethod
-    def attention_backend(self, use_mla: bool = False) -> str:
+    def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> str:
         """
         Get the attention backend class path for this platform.
 
@@ -138,6 +138,7 @@ def attention_backend(self, use_mla: bool = False) -> str:
 
         Args:
             use_mla: Whether to use Multi-head Latent Attention (MLA)
+            use_sparse: Whether to use sparse attention
 
         Returns:
             Fully qualified class path string, e.g.:

vllm_fl/platform.py

@@ -170,12 +170,14 @@ def get_attn_backend_cls(
         from vllm_fl.dispatch import call_op
 
         use_mla = attn_selector_config.use_mla
+        use_sparse = attn_selector_config.use_sparse
 
-        backend_path = call_op("attention_backend", use_mla=use_mla)
+        backend_path = call_op("attention_backend", use_mla=use_mla, use_sparse=use_sparse)
 
         logger.info_once(
-            "Using attention backend via dispatch (use_mla=%s): %s",
+            "Using attention backend via dispatch (use_mla=%s, use_sparse=%s): %s",
             use_mla,
+            use_sparse,
             backend_path,
             scope="local",
         )