From b2b02108736d30ab1c3b3d2199ce33e318df3201 Mon Sep 17 00:00:00 2001 From: "(Messi) Junlin Wu" Date: Wed, 29 Jul 2026 20:01:04 +0800 Subject: [PATCH] :building_construction: refactor(platform): add typed device features --- python/sglang/srt/platforms/capabilities.py | 12 ++++ python/sglang/srt/platforms/device_mixin.py | 9 +++ python/sglang/srt/platforms/npu.py | 59 +++++++++++++++++ .../unit/platforms/test_platform_interface.py | 63 +++++++++++++++++++ 4 files changed, 143 insertions(+) create mode 100644 python/sglang/srt/platforms/capabilities.py diff --git a/python/sglang/srt/platforms/capabilities.py b/python/sglang/srt/platforms/capabilities.py new file mode 100644 index 000000000000..557db561a663 --- /dev/null +++ b/python/sglang/srt/platforms/capabilities.py @@ -0,0 +1,12 @@ +"""Semantic hardware features exposed by SGLang platforms.""" + +from enum import Enum + + +class PlatformFeature(str, Enum): + """Optional hardware/runtime features queried through ``current_platform``.""" + + NPU_NATIVE_GEMMA_RMS_NORM = "npu.native_gemma_rms_norm" + + +__all__ = ["PlatformFeature"] diff --git a/python/sglang/srt/platforms/device_mixin.py b/python/sglang/srt/platforms/device_mixin.py index a87523aa72f6..bb01cd1f4d1f 100644 --- a/python/sglang/srt/platforms/device_mixin.py +++ b/python/sglang/srt/platforms/device_mixin.py @@ -33,6 +33,7 @@ import torch from sglang.srt.environ import envs +from sglang.srt.platforms.capabilities import PlatformFeature class PlatformEnum(enum.Enum): @@ -159,6 +160,14 @@ def get_current_memory_usage( """[Active] Get current peak memory usage in bytes.""" raise NotImplementedError + def supports_feature(self, feature: PlatformFeature, device_id: int = 0) -> bool: + """[Active] Whether this platform supports an optional semantic feature. + + Platforms opt in explicitly. The conservative default keeps unknown and + out-of-tree devices on the caller's safe fallback implementation. + """ + return False + # ------------------------------------------------------------------ # Planned methods — reserved interface. Core still uses hardcoded # calls (e.g. torch.cuda.*). OOT implementations will NOT take diff --git a/python/sglang/srt/platforms/npu.py b/python/sglang/srt/platforms/npu.py index 44a34236e910..f946d3f8e2a9 100644 --- a/python/sglang/srt/platforms/npu.py +++ b/python/sglang/srt/platforms/npu.py @@ -13,11 +13,15 @@ ``device == "npu"`` string checks. """ +import logging +from enum import Enum +from functools import cached_property from typing import Optional import torch from sglang.srt.environ import envs +from sglang.srt.platforms.capabilities import PlatformFeature from sglang.srt.platforms.device_mixin import ( DeviceCapability, DeviceMixin, @@ -25,6 +29,31 @@ ) from sglang.srt.platforms.interface import SRTPlatform +logger = logging.getLogger(__name__) + + +class NpuSocFamily(str, Enum): + """Ascend product families used only to build the platform feature set.""" + + A2 = "a2" + A3 = "a3" + A5 = "a5" + UNKNOWN = "unknown" + + +_NPU_SOC_FAMILY_RANGES = ( + (220, 225, NpuSocFamily.A2), + (250, 255, NpuSocFamily.A3), + (260, 260, NpuSocFamily.A5), +) + +_NPU_FEATURES_BY_SOC = { + NpuSocFamily.A2: frozenset({PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM}), + NpuSocFamily.A3: frozenset({PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM}), + NpuSocFamily.A5: frozenset(), + NpuSocFamily.UNKNOWN: frozenset(), +} + class NpuDeviceMixin(DeviceMixin): """Ascend NPU implementation of the shared device operations.""" @@ -33,6 +62,36 @@ class NpuDeviceMixin(DeviceMixin): device_name: str = "npu" device_type: str = "npu" + @cached_property + def soc_family(self) -> NpuSocFamily: + """Resolve the process-wide Ascend SoC family once, on first use.""" + try: + import torch_npu + + soc_version = int(torch_npu.npu.get_soc_version()) + except (ImportError, AttributeError, RuntimeError, TypeError, ValueError) as e: + logger.warning( + "Unable to resolve the Ascend SoC version; optional platform " + "features will use safe fallbacks: %s", + e, + ) + return NpuSocFamily.UNKNOWN + + for lower, upper, family in _NPU_SOC_FAMILY_RANGES: + if lower <= soc_version <= upper: + return family + + logger.warning( + "Unknown Ascend SoC version %s; optional platform features will " + "use safe fallbacks.", + soc_version, + ) + return NpuSocFamily.UNKNOWN + + def supports_feature(self, feature: PlatformFeature, device_id: int = 0) -> bool: + # torch_npu exposes SoC version process-wide rather than per device. + return feature in _NPU_FEATURES_BY_SOC[self.soc_family] + def get_device_total_memory(self, device_id: int = 0) -> int: return int(torch.npu.mem_get_info(device_id)[1]) diff --git a/test/registered/unit/platforms/test_platform_interface.py b/test/registered/unit/platforms/test_platform_interface.py index 2dcc2efbfbd7..4986dba87a04 100644 --- a/test/registered/unit/platforms/test_platform_interface.py +++ b/test/registered/unit/platforms/test_platform_interface.py @@ -5,11 +5,14 @@ and the platform discovery / lazy initialization mechanism. """ +import sys +from types import SimpleNamespace from unittest.mock import MagicMock, patch import torch from sglang.srt.platforms import _load_platform_class, _resolve_platform +from sglang.srt.platforms.capabilities import PlatformFeature from sglang.srt.platforms.cpu import CpuDeviceMixin, CpuSRTPlatform from sglang.srt.platforms.cuda import CudaDeviceMixin, CudaSRTPlatform from sglang.srt.platforms.device_mixin import ( @@ -19,6 +22,7 @@ PlatformEnum, ) from sglang.srt.platforms.interface import SRTPlatform +from sglang.srt.platforms.npu import NpuSocFamily, NpuSRTPlatform from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -194,6 +198,12 @@ def test_is_out_of_tree(self): cuda = _make_device_mixin(PlatformEnum.CUDA, "cuda", "cuda") self.assertFalse(cuda.is_out_of_tree()) + def test_optional_features_default_to_false(self): + mixin = _make_device_mixin(PlatformEnum.OOT, "custom", "custom") + self.assertFalse( + mixin.supports_feature(PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM) + ) + @patch("platform.machine") def test_get_cpu_architecture(self, mock_machine): """get_cpu_architecture maps common strings to CpuArchEnum.""" @@ -449,6 +459,59 @@ def test_cpu_srt_platform_capabilities(self): self.assertFalse(base.is_pin_memory_available()) +class TestNpuDeviceMixin(CustomTestCase): + def _platform_for_soc(self, soc_version): + get_soc_version = MagicMock(return_value=soc_version) + torch_npu = SimpleNamespace( + npu=SimpleNamespace(get_soc_version=get_soc_version) + ) + return NpuSRTPlatform(), torch_npu, get_soc_version + + def test_native_gemma_rms_norm_feature_matrix(self): + cases = [ + (220, NpuSocFamily.A2, True), + (225, NpuSocFamily.A2, True), + (250, NpuSocFamily.A3, True), + (255, NpuSocFamily.A3, True), + (260, NpuSocFamily.A5, False), + (999, NpuSocFamily.UNKNOWN, False), + ] + for soc_version, expected_family, expected_support in cases: + with self.subTest(soc_version=soc_version): + platform, torch_npu, _ = self._platform_for_soc(soc_version) + with patch.dict(sys.modules, {"torch_npu": torch_npu}): + self.assertEqual(platform.soc_family, expected_family) + self.assertEqual( + platform.supports_feature( + PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM + ), + expected_support, + ) + + def test_soc_detection_is_cached(self): + platform, torch_npu, get_soc_version = self._platform_for_soc(220) + with patch.dict(sys.modules, {"torch_npu": torch_npu}): + self.assertTrue( + platform.supports_feature(PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM) + ) + self.assertTrue( + platform.supports_feature(PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM) + ) + get_soc_version.assert_called_once_with() + + def test_soc_detection_failure_uses_safe_fallback(self): + get_soc_version = MagicMock(side_effect=RuntimeError("not initialized")) + torch_npu = SimpleNamespace( + npu=SimpleNamespace(get_soc_version=get_soc_version) + ) + platform = NpuSRTPlatform() + with patch.dict(sys.modules, {"torch_npu": torch_npu}): + self.assertEqual(platform.soc_family, NpuSocFamily.UNKNOWN) + self.assertFalse( + platform.supports_feature(PlatformFeature.NPU_NATIVE_GEMMA_RMS_NORM) + ) + + class TestSRTPlatformOverrides(CustomTestCase): """Tests for SRTPlatform method overrides via plugins."""