fix(ascend): detect exact SoC version

src/native/ascend/custom/cmake/detect_soc.cmake

@@ -1,20 +1,64 @@
-# Auto-detect the Ascend SOC version from `npu-smi info`.
+# Auto-detect the Ascend SOC version from torch_npu or `npu-smi info`.
 #
-# `infiniops_detect_soc(<out_var>)` parses the first `910*` / `310*` entry
-# in `npu-smi info` and writes `Ascend<NNNX>` into the named variable in
-# the caller's scope.  Falls back to `Ascend910B4` when detection fails
-# (no NPU on the host, `npu-smi` missing, output format mismatch).
+# `infiniops_detect_soc(<out_var>)` first asks torch_npu for the exact device
+# name, then falls back to a model-qualified `910*` / `310*` token from
+# `npu-smi info`.  Falls back to `Ascend910B4` when detection fails (no NPU
+# on the host, missing tools, or an output format without a model suffix).
 #
 # Called from both `src/CMakeLists.txt` (outer `pip install` build, to
 # forward `SOC_VERSION` to the standalone `build.sh` invocation) and
 # `src/native/ascend/custom/cmake/config_ascend.cmake` (the sub-build driven
 # by that `build.sh`).
 
 function(infiniops_detect_soc out_var)
-    execute_process(
-        COMMAND bash -c "npu-smi info 2>/dev/null | awk '/910B|910A|310/ {for (i=1;i<=NF;i++) if ($i ~ /^(910|310)/) {print \"Ascend\" $i; exit}}'"
-        OUTPUT_VARIABLE _detected
-        OUTPUT_STRIP_TRAILING_WHITESPACE)
+    set(_detected "")
+    set(_python "")
+
+    if(DEFINED _TORCH_PYTHON AND NOT "${_TORCH_PYTHON}" STREQUAL "")
+        set(_python "${_TORCH_PYTHON}")
+    elseif(DEFINED Python_EXECUTABLE AND NOT "${Python_EXECUTABLE}" STREQUAL "")
+        set(_python "${Python_EXECUTABLE}")
+    else()
+        find_program(_python NAMES python3 python)
+    endif()
+
+    if(_python)
+        execute_process(
+            COMMAND "${_python}" -c "import torch, torch_npu; print(torch.npu.get_device_name(0))"
+            RESULT_VARIABLE _torch_npu_result
+            OUTPUT_VARIABLE _torch_npu_output
+            ERROR_QUIET
+            OUTPUT_STRIP_TRAILING_WHITESPACE)
+
+        if(_torch_npu_result EQUAL 0 AND
+           _torch_npu_output MATCHES "^Ascend(910|310)[A-Za-z0-9_]*$")
+            set(_detected "${_torch_npu_output}")
+        endif()
+    endif()
+
+    if(NOT _detected)
+        execute_process(
+            COMMAND npu-smi info
+            RESULT_VARIABLE _npu_smi_result
+            OUTPUT_VARIABLE _npu_smi_output
+            ERROR_QUIET
+            OUTPUT_STRIP_TRAILING_WHITESPACE)
+
+        if(_npu_smi_result EQUAL 0)
+            string(REGEX MATCH
+                "(Ascend)?(910|310)[A-Za-z_][A-Za-z0-9_]*"
+                _npu_smi_soc
+                "${_npu_smi_output}")
+
+            if(_npu_smi_soc)
+                if(_npu_smi_soc MATCHES "^Ascend")
+                    set(_detected "${_npu_smi_soc}")
+                else()
+                    set(_detected "Ascend${_npu_smi_soc}")
+                endif()
+            endif()
+        endif()
+    endif()
 
     if(_detected)
         set(${out_var} "${_detected}" PARENT_SCOPE)

tests/test_ascend_soc_detection.py

@@ -0,0 +1,90 @@
+import os
+import stat
+import subprocess
+from pathlib import Path
+
+import pytest
+
+
+DETECT_SOC = (
+    Path(__file__).parents[1]
+    / "src"
+    / "native"
+    / "ascend"
+    / "custom"
+    / "cmake"
+    / "detect_soc.cmake"
+)
+
+
+def _write_executable(path, contents):
+    path.write_text(contents)
+    path.chmod(path.stat().st_mode | stat.S_IXUSR)
+
+
+def _detect_soc(tmp_path, python_output=None, npu_smi_output=""):
+    if os.name == "nt":
+        pytest.skip("the detector invokes POSIX Ascend command-line tools")
+
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    python = bin_dir / "python3"
+    npu_smi = bin_dir / "npu-smi"
+
+    if python_output is None:
+        _write_executable(python, "#!/bin/sh\nexit 1\n")
+    else:
+        _write_executable(python, f"#!/bin/sh\nprintf '%s\\n' '{python_output}'\n")
+
+    _write_executable(
+        npu_smi,
+        "#!/bin/sh\ncat <<'EOF'\n" + npu_smi_output + "\nEOF\n",
+    )
+
+    driver = tmp_path / "detect.cmake"
+    driver.write_text(
+        f'set(Python_EXECUTABLE "{python.as_posix()}")\n'
+        f'include("{DETECT_SOC.as_posix()}")\n'
+        "infiniops_detect_soc(DETECTED_SOC)\n"
+        'message(STATUS "DETECTED_SOC=${DETECTED_SOC}")\n'
+    )
+    env = os.environ.copy()
+    env["PATH"] = f"{bin_dir}{os.pathsep}{env['PATH']}"
+    result = subprocess.run(
+        ["cmake", "-P", str(driver)],
+        capture_output=True,
+        text=True,
+        env=env,
+        check=True,
+    )
+
+    return result.stdout
+
+
+def test_detect_soc_prefers_exact_torch_npu_name(tmp_path):
+    output = _detect_soc(
+        tmp_path,
+        python_output="Ascend910_9362",
+        npu_smi_output="| 7 Ascend910 | 0 0 / 0 3109 / 65536 |",
+    )
+
+    assert "DETECTED_SOC=Ascend910_9362" in output
+
+
+def test_detect_soc_ignores_hbm_numbers_in_npu_smi(tmp_path):
+    output = _detect_soc(
+        tmp_path,
+        npu_smi_output="| 7 Ascend910 | 0 0 / 0 3109 / 65536 |",
+    )
+
+    assert "DETECTED_SOC=Ascend910B4" in output
+    assert "Ascend3109" not in output
+
+
+def test_detect_soc_accepts_model_qualified_npu_smi_name(tmp_path):
+    output = _detect_soc(
+        tmp_path,
+        npu_smi_output="| 0 910B4 | 0 0 / 0 2789 / 32768 |",
+    )
+
+    assert "DETECTED_SOC=Ascend910B4" in output