[NPU] NZ for non-quantized MOE, Qwen3 MOE double memory consumption fix (#15904)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-01-29 00:55:08 +08:00
parent 1953efb60e
commit b77b0ffd60
7 changed files with 111 additions and 49 deletions
@@ -0,0 +1,76 @@
+"""
+Usage:
+python3 -m unittest test_ascend_memory_consumption.TestMemoryConsumptionAscend.test_memory_consumption
+"""
+
+import os
+import unittest
+
+import torch
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_npu_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_npu_ci(est_time=400, suite="nightly-1-npu-a3", nightly=True)
+
+if "ASCEND_RT_VISIBLE_DEVICES" not in os.environ:
+    os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "0,1"
+DEFAULT_PORT_FOR_SRT_TEST_RUNNER = (
+    8000 + int(os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "0")[0]) * 100
+)
+DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
+
+
+class TestMemoryConsumptionAscend(CustomTestCase):
+
+    def test_memory_consumption(self):
+
+        model = "nytopop/Qwen3-30B-A3B.w8a8"
+        base_url = DEFAULT_URL_FOR_TEST
+
+        ### Calculate initial used memory
+        free_npu_memory, total_npu_memory = torch.npu.mem_get_info()
+        initial_used_memory = total_npu_memory - free_npu_memory
+
+        process = popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--device",
+                "npu",
+                "--attention-backend",
+                "ascend",
+                "--tp-size",
+                "2",
+                "--mem-fraction-static",
+                "0.8",
+                "--cuda-graph-bs",
+                "1",
+                "--max-total-tokens",
+                "1024",
+                "--disable-radix-cache",
+                "--disable-cuda-graph",
+            ],
+        )
+
+        ### Calculate initial used memory
+        free_npu_memory, total_npu_memory = torch.npu.mem_get_info()
+        used_memory_after_server_starting = (
+            total_npu_memory - free_npu_memory - initial_used_memory
+        ) / (1 << 30)
+        self.assertLessEqual(float(used_memory_after_server_starting), 16.00)
+
+        # Clean up everything
+        kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()