[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:
Chunyu
2026-03-18 10:21:31 +00:00
committed by GitHub
parent 05f5878201
commit f762d98202
@@ -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")