Detect CUDA UMD Version from newer nvidia-smi output (fixes #5812) (#5817)

* Detect CUDA UMD Version from newer nvidia-smi output (#5812)

Newer NVIDIA drivers (e.g. 610.x on Windows) print the driver's max
CUDA capability as "CUDA UMD Version: X.Y" instead of the legacy
"CUDA Version: X.Y" header.  The installers and Studio setup scripts
were only matching the legacy spelling, so on a fresh RTX 5090
laptop with a 13.x driver they failed to detect any CUDA version
and fell through to the cu126 wheel default.

Accept both spellings everywhere we parse nvidia-smi output:

- install.ps1: Get-TorchIndexUrl regex now allows " UMD"
- install.sh: two-expression sed (POSIX BRE has no "?"); the two
  patterns are mutually exclusive per line, head -1 picks the match
- studio/setup.ps1: Get-PytorchCudaTag and the $DriverMaxCuda
  detector both relaxed
- studio/install_llama_prebuilt.py: substring scan replaced with a
  regex search using the same pattern
- tests/sh/test_get_torch_index_url.sh: new make_mock_smi_umd helper
  plus three UMD cases (13.3 -> cu130, 12.8 -> cu128, 11.8 -> cu118);
  all 30 tests pass locally

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-05-27 10:37:21 -07:00 committed by GitHub
commit a62eb80f7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 73 additions and 11 deletions

View file

@ -1253,7 +1253,10 @@ shell.Run cmd, 0, False
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = & $NvidiaSmiExe 2>&1 | Out-String
if ($output -match 'CUDA Version:\s+(\d+)\.(\d+)') {
# Newer NVIDIA drivers (e.g. 610.x on Windows) print
# "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y".
# Accept both spellings so we don't fall through to the cu126 default.
if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') {
$major = [int]$Matches[1]; $minor = [int]$Matches[2]
if ($major -ge 13) { return "$baseUrl/cu130" }
if ($major -eq 12 -and $minor -ge 8) { return "$baseUrl/cu128" }