[Feature] Add MLAProcess for DeepSeek MLA on NPU (#10130)

2025-09-23 08:17:48 +08:00
parent 0753ef831e
commit d27a6f7092
7 changed files with 369 additions and 23 deletions
--- a/python/sglang/srt/utils.py
+++ b/python/sglang/srt/utils.py
@@ -174,6 +174,8 @@ def is_blackwell():

@lru_cache(maxsize=1)
 def is_sm100_supported(device=None) -> bool:
+    if not is_cuda_alike():
+        return False
    return (torch.cuda.get_device_capability(device)[0] == 10) and (
        torch.version.cuda >= "12.8"
    )
@@ -181,6 +183,8 @@ def is_sm100_supported(device=None) -> bool:

@lru_cache(maxsize=1)
 def is_sm90_supported(device=None) -> bool:
+    if not is_cuda_alike():
+        return False
    return (torch.cuda.get_device_capability(device)[0] == 9) and (
        torch.version.cuda >= "12.3"
    )