[NPU] Add sparsity-driven KV offload for DeepSeek DSA on Ascend (#33089)

Co-authored-by: itcj <itcj@foxmail.com> Co-authored-by: chenyh <chenyh18@mail.ustc.edu.cn> Co-authored-by: hibikid <7514545+hibikid@user.noreply.gitee.com> Co-authored-by: Chengjie Tang <tangcj@sxu.edu.cn> Co-authored-by: wangbiao0814 <wangbiao0814@outlook.com>
2026-09-09 09:13:22 +08:00
parent a8e45f16cc
commit 295132c4a5
9 changed files with 1737 additions and 30 deletions
@@ -0,0 +1,121 @@
+import os
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config import (
+    get_sparsity_driven_kv_offload_cell_size,
+    get_sparsity_driven_kv_offload_sparse_context_len,
+    is_sparsity_driven_kv_offload_enabled,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+def _make_glm51_model_config():
+    hf_config = SimpleNamespace(
+        architectures=["GlmMoeDsaForCausalLM"],
+        index_head_dim=128,
+        index_topk=1536,
+    )
+    hf_config.get_text_config = lambda: hf_config
+    return SimpleNamespace(
+        hf_config=hf_config,
+        index_head_dim=128,
+    )
+
+
+class TestSparsityDrivenKVOffloadConfig(unittest.TestCase):
+    def test_glm_dsa_model_enables_sparse_kv_offload(self):
+        with (
+            patch.dict(
+                os.environ,
+                {"SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD": "1"},
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.is_npu",
+                return_value=True,
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.attention_backends",
+                return_value=("ascend", "ascend"),
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.get_schedule",
+                return_value=SimpleNamespace(max_running_requests=8),
+            ),
+        ):
+            model_config = _make_glm51_model_config()
+
+            self.assertTrue(
+                is_sparsity_driven_kv_offload_enabled(
+                    model_config=model_config,
+                    use_mla_backend=True,
+                )
+            )
+            self.assertEqual(
+                get_sparsity_driven_kv_offload_sparse_context_len(
+                    model_config=model_config
+                ),
+                1536,
+            )
+            self.assertEqual(
+                get_sparsity_driven_kv_offload_cell_size(
+                    model_config=model_config,
+                    use_mla_backend=True,
+                    num_layers=2,
+                    element_size=2,
+                ),
+                512,
+            )
+
+    def test_split_attention_backend_rejects_sparse_kv_offload(self):
+        with (
+            patch.dict(
+                os.environ,
+                {"SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD": "1"},
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.is_npu",
+                return_value=True,
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.attention_backends",
+                return_value=("ascend", "torch_native"),
+            ),
+        ):
+            with self.assertRaisesRegex(ValueError, "Ascend MLA attention backend"):
+                is_sparsity_driven_kv_offload_enabled(
+                    model_config=_make_glm51_model_config(),
+                    use_mla_backend=True,
+                )
+
+    def test_missing_request_capacity_rejects_sparse_kv_offload(self):
+        with (
+            patch.dict(
+                os.environ,
+                {"SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD": "1"},
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.is_npu",
+                return_value=True,
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.attention_backends",
+                return_value=("ascend", "ascend"),
+            ),
+            patch(
+                "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.get_schedule",
+                return_value=SimpleNamespace(max_running_requests=None),
+            ),
+        ):
+            with self.assertRaisesRegex(ValueError, "max_running_requests"):
+                is_sparsity_driven_kv_offload_enabled(
+                    model_config=_make_glm51_model_config(),
+                    use_mla_backend=True,
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()