[NPU]Support torch_npu profiler patch API drift (#26356)

Co-authored-by: leland17 <lileliao@foxmail.com>
Co-authored-by: OmX <omx@oh-my-codex.dev>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/multimodal_gen/runtime/utils/profiler.py

@@ -5,6 +5,7 @@ import torch
 
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
 
 if current_platform.is_npu():
     import torch_npu
@@ -13,7 +14,7 @@ if current_platform.is_npu():
         ["profiler.profile", torch_npu.profiler.profile],
         ["profiler.schedule", torch_npu.profiler.schedule],
     ]
-    torch_npu._apply_patches(patches)
+    apply_torch_npu_patches(torch_npu, patches)
 
 logger = init_logger(__name__)
 

python/sglang/srt/managers/scheduler_components/profiler_manager.py

@@ -21,6 +21,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import is_npu
 from sglang.srt.utils.profile_merger import ProfileMerger
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -34,7 +35,7 @@ if _is_npu:
         ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU],
         ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU],
     ]
-    torch_npu._apply_patches(patches)
+    apply_torch_npu_patches(torch_npu, patches)
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/utils/profile_utils.py

@@ -13,6 +13,7 @@ from sglang.srt.managers.io_struct import ProfileReqOutput
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import is_npu
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
 
 _is_npu = is_npu()
 if _is_npu:
@@ -23,7 +24,7 @@ if _is_npu:
         ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU],
         ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU],
     ]
-    torch_npu._apply_patches(patches)
+    apply_torch_npu_patches(torch_npu, patches)
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/utils/torch_npu_patch_utils.py

@@ -0,0 +1,17 @@
+from collections.abc import Sequence
+from typing import Any
+
+
+def apply_torch_npu_patches(torch_npu: Any, patches: Sequence[Sequence[Any]]) -> None:
+    """Apply torch_npu patches across old and new torch_npu patch APIs."""
+    if hasattr(torch_npu, "_apply_patches"):
+        torch_npu._apply_patches(patches)
+        return
+
+    if hasattr(torch_npu, "_apply_all_patches"):
+        torch_npu._apply_all_patches()
+        return
+
+    raise AttributeError(
+        "torch_npu must provide either _apply_patches or _apply_all_patches"
+    )

test/registered/unit/utils/test_torch_npu_patch_utils.py

@@ -0,0 +1,41 @@
+import types
+import unittest
+
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestTorchNpuPatchUtils(unittest.TestCase):
+    def test_apply_torch_npu_patches_uses_targeted_api_when_available(self):
+        calls = []
+        torch_npu = types.SimpleNamespace(
+            _apply_patches=lambda patches: calls.append(("_apply_patches", patches)),
+            _apply_all_patches=lambda: calls.append(("_apply_all_patches", None)),
+        )
+        patches = [["profiler.profile", object()]]
+
+        apply_torch_npu_patches(torch_npu, patches)
+
+        self.assertEqual(calls, [("_apply_patches", patches)])
+
+    def test_apply_torch_npu_patches_uses_all_patches_api_when_targeted_api_missing(
+        self,
+    ):
+        calls = []
+        torch_npu = types.SimpleNamespace(
+            _apply_all_patches=lambda: calls.append("_apply_all_patches")
+        )
+
+        apply_torch_npu_patches(torch_npu, [["profiler.profile", object()]])
+
+        self.assertEqual(calls, ["_apply_all_patches"])
+
+    def test_apply_torch_npu_patches_requires_supported_api(self):
+        with self.assertRaises(AttributeError):
+            apply_torch_npu_patches(types.SimpleNamespace(), [])
+
+
+if __name__ == "__main__":
+    unittest.main()