mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
[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
This commit is contained in:
@@ -198,6 +198,12 @@ class UBManager:
|
||||
if is_npu_available():
|
||||
try:
|
||||
from tbe.common.platform import get_soc_spec
|
||||
from tbe.common.platform import set_current_compile_soc_info
|
||||
|
||||
# Set current SOC info for get_soc_spec to work correctly
|
||||
device = getattr(torch, "npu")
|
||||
soc_info = device.get_device_name(device.current_device())
|
||||
set_current_compile_soc_info(soc_info)
|
||||
|
||||
# Query UB size (get_soc_spec returns size in bytes)
|
||||
ub_size_bytes = get_soc_spec("UB_SIZE")
|
||||
|
||||
Reference in New Issue
Block a user