Skip to content

Commit f762d98

Browse files
authored
[Fix, NPU] Set soc_info for the NPU device (#1149)
## Summary With the `teb` update, the `soc version` must be explicitly set in `set_current_compile_soc_info` for `get_soc_spec` to return the correct `UB_SIZE`. ``` import torch from tbe.common.platform import get_soc_spec from tbe.common.platform import set_current_compile_soc_info ub_size = get_soc_spec("UB_SIZE") print(f"UB size: {ub_size}") ``` <img width="646" height="32" alt="image" src="https://github.com/user-attachments/assets/d57e52e5-e468-4e46-a4db-94351364d218" /> ❌ ## Testing Done ### Atlas 900 A2 PoD #### Script: ub_size_test.py ```python import torch from tbe.common.platform import get_soc_spec from tbe.common.platform import set_current_compile_soc_info device = getattr(torch, "npu") soc_info = device.get_device_name(device.current_device()) print("Current SOC info:", soc_info) ub_size = get_soc_spec("UB_SIZE") print(f"UB size: {ub_size}") ``` #### Result <img width="713" height="304" alt="image" src="https://github.com/user-attachments/assets/56215e32-1c1b-49ed-9e5b-a859e80729d1" /> ### Atlas 800I A2 Thanks @TianHao324 for verifying on the Atlas 800I A2 — results are correct there too. - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence
1 parent 05f5878 commit f762d98

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

src/liger_kernel/ops/backends/_ascend/ub_manager.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,12 @@ def _detect_ub_capacity(self) -> int:
198198
if is_npu_available():
199199
try:
200200
from tbe.common.platform import get_soc_spec
201+
from tbe.common.platform import set_current_compile_soc_info
202+
203+
# Set current SOC info for get_soc_spec to work correctly
204+
device = getattr(torch, "npu")
205+
soc_info = device.get_device_name(device.current_device())
206+
set_current_compile_soc_info(soc_info)
201207

202208
# Query UB size (get_soc_spec returns size in bytes)
203209
ub_size_bytes = get_soc_spec("UB_SIZE")

0 commit comments

Comments
 (0)