diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index ed76d67b8..75b3cc138 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/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__) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 6e5a58ffa..31df519f9 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/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__) diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 0b2fbeccc..7b6bcac31 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/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__) diff --git a/python/sglang/srt/utils/torch_npu_patch_utils.py b/python/sglang/srt/utils/torch_npu_patch_utils.py new file mode 100644 index 000000000..94dc0a634 --- /dev/null +++ b/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" + ) diff --git a/test/registered/unit/utils/test_torch_npu_patch_utils.py b/test/registered/unit/utils/test_torch_npu_patch_utils.py new file mode 100644 index 000000000..09f08df25 --- /dev/null +++ b/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()