From 34f35469294f536da871be416ee93c35f1f5db98 Mon Sep 17 00:00:00 2001 From: Even Zhou Date: Fri, 28 Nov 2025 15:41:46 +0800 Subject: [PATCH 1/2] update utils --- python/sgl_kernel_npu/sgl_kernel_npu/utils.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/utils.py diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/utils.py b/python/sgl_kernel_npu/sgl_kernel_npu/utils.py new file mode 100644 index 00000000..12b05840 --- /dev/null +++ b/python/sgl_kernel_npu/sgl_kernel_npu/utils.py @@ -0,0 +1,43 @@ +from enum import IntEnum +from functools import lru_cache +from typing import Tuple + +import torch +import torch_npu + + +class DeviceCapability(IntEnum): + """ + Refer to `SocVersion` in https://gitcode.com/Ascend/pytorch/blob/master/torch_npu/csrc/core/npu/NpuVariables.h + + Please be aware that this is an internal interface which is subjected to change without prior notice. + """ + + # Ascend 910B1,910B2,910B2C,910B3,910B4,910B4_1 + A2 = 220 + + # Ascend 910_9391,910_9392,910_9381,910_9382,910_9372,910_9362, + A3 = 250 + + +@lru_cache(maxsize=1) +def get_device_properties() -> Tuple[int, int]: + device = torch.npu.current_device() + device_properties = torch.npu.get_device_properties(device) + + cube_core_num = device_properties.cube_core_num + vector_core_num = device_properties.vector_core_num + + return cube_core_num, vector_core_num + + +@lru_cache(maxsize=1) +def get_device_capability() -> DeviceCapability: + soc = torch_npu._C._npu_get_soc_version() + soc = soc // 10 * 10 + + assert soc in iter( + DeviceCapability + ), f"Unsupported soc: {torch.npu.get_device_name()}" + + return DeviceCapability(soc) From ee4425eefeb1268e52b81952f6647942d06f2e16 Mon Sep 17 00:00:00 2001 From: Even Zhou Date: Fri, 28 Nov 2025 15:50:19 +0800 Subject: [PATCH 2/2] update utils --- python/sgl_kernel_npu/sgl_kernel_npu/utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/utils.py b/python/sgl_kernel_npu/sgl_kernel_npu/utils.py index 12b05840..c70911d4 100644 --- a/python/sgl_kernel_npu/sgl_kernel_npu/utils.py +++ b/python/sgl_kernel_npu/sgl_kernel_npu/utils.py @@ -20,6 +20,12 @@ class DeviceCapability(IntEnum): A3 = 250 +@lru_cache(maxsize=1) +def get_device_name() -> str: + device = torch.npu.current_device() + return torch.npu.get_device_name(device) + + @lru_cache(maxsize=1) def get_device_properties() -> Tuple[int, int]: device = torch.npu.current_device() @@ -36,8 +42,6 @@ def get_device_capability() -> DeviceCapability: soc = torch_npu._C._npu_get_soc_version() soc = soc // 10 * 10 - assert soc in iter( - DeviceCapability - ), f"Unsupported soc: {torch.npu.get_device_name()}" + assert soc in iter(DeviceCapability), f"Unsupported soc: {get_device_name()}" return DeviceCapability(soc)