uv-torch 深度解析：uv 如何自动选择 PyTorch 的 CUDA、ROCm 与 XPU 索引

本文围绕 uv 仓库中的 crates/uv-torch 组件展开，讲清 uv 在安装 PyTorch 生态包时如何根据操作系统、GPU 驱动版本自动或手动选择正确的 PyTorch 索引（--torch-backend 选项），并深入剖析其硬件探测流程（CUDA 驱动版本读取、AMD GPU 架构识别、Intel XPU 检测）与后端到索引 URL 的映射机制。读完后，你既能正确使用 uv --torch-backend 完成 GPU 版 PyTorch 安装，也能从源码层面理解自动探测（auto 模式）背后的完整决策链。

一、uv-torch 是什么：为 PyTorch 生态选对索引的底层组件

crates/uv-torch/README.md 明确说明：uv-torch 是 uv 的内部组件 crate（当前版本 0.0.76，对应 uv 0.12.9），其对外暴露的 Rust API 不稳定，会频繁发生破坏性变更。它的核心职责在 crates/uv-torch/src/backend.rs 的模块注释中一句话概括：

"uv-torch is a library for determining the appropriate PyTorch index based on the operating system and CUDA driver version."

即：根据操作系统与加速硬件状态，确定应使用哪个 PyTorch 官方索引。PyTorch 官方将 CPU、各 CUDA 版本、各 ROCm 版本以及 XPU 的轮子分别发布在不同的索引 URL 下（如 https://download.pytorch.org/whl/cpu、https://download.pytorch.org/whl/cu128），如果索引选错，用户要么装到 CPU 版 torch 导致 GPU 不可用，要么装到与本机驱动不兼容的 CUDA 构建。uv-torch 就是解决这一问题的"索引路由器"。

从 crates/uv-torch/src/lib.rs 可以看到 crate 结构非常精炼，仅两个模块：

mod accelerator;
mod backend;

pub use accelerator::{AcceleratorError, AmdGpuArchitecture};
pub use backend::*;

• accelerator.rs：负责硬件加速器的探测（NVIDIA/AMD/Intel）；
• backend.rs：负责后端（backend）的定义、策略计算与索引 URL 映射。

crates/uv-torch/Cargo.toml 显示其依赖了 uv-distribution-types（IndexUrl 类型）、uv-normalize（包名）、uv-pep440（版本号）、uv-platform-tags（操作系统标识），并在 Windows 目标下额外依赖 windows crate 用于查询设备树。

值得一提的是，backend.rs 文件头的许可证注释（backend.rs#L4-L37）声明该库的探测逻辑派生自 Philipp Meier 的 light-the-torch 项目（BSD-3-Clause 授权）——这正是早期 pip install light-the-torch && light-the-torch 手动修补 pip 的方案在 uv 中以 Rust 原生化实现的版本。

二、用户入口：--torch-backend 选项与 TorchMode

2.1 CLI 上的使用方式

在 crates/uv-cli/src/lib.rs#L1919-L1932 中，--torch-backend 参数的官方文档注释给出了完整用法说明：

The backend to use when fetching packages in the PyTorch ecosystem (e.g., `cpu`, `cu126`, or `auto`).

When set, uv will ignore the configured index URLs for packages in the PyTorch ecosystem,
and will instead use the defined backend.

For example, when set to `cpu`, uv will use the CPU-only PyTorch index; when set to `cu126`,
uv will use the PyTorch index for CUDA 12.6.

The `auto` mode will attempt to detect the appropriate PyTorch index based on the currently
installed CUDA drivers.

This option is in preview and may change in any future release.

要点：

1. 作用域：设置后，uv 会忽略你为 PyTorch 生态包配置的索引 URL，改用该 backend 对应的官方索引；
2. 取值：auto、cpu、cu80～cu132、rocm4.0.1～rocm7.2、xpu；
3. 预览特性：该选项处于 preview 状态，后续版本可能调整；
4. 环境变量：#[arg(long, value_enum, env = EnvVars::UV_TORCH_BACKEND)] 表明其还可通过环境变量 UV_TORCH_BACKEND 设置。

从 crates/uv-cli/src/lib.rs 的定义看，torch_backend: Option<TorchMode> 字段出现在 pip compile、pip install、pip sync、uv add 等多个命令参数结构中（如 L1932、L2209、L2561、L5709、L5902），并且 crates/uv/src/settings.rs 显示它会按"命令行参数 > 项目配置 > 全局配置"的优先级合并（args.torch_backend.combine(torch_backend)），最终落入 ResolverSettings 供解析阶段使用。

2.2 TorchMode：用户可见的全部取值

TorchMode 枚举 定义了所有可用模式，按序列化规则 rename_all = "kebab-case" 及个别 #[serde(rename = ...)] 特例，对应 CLI/配置中的字面量：

类别	取值（CLI 字面量）	对应索引
自动	auto	根据系统检测
CPU	cpu	https://download.pytorch.org/whl/cpu
CUDA	cu80, cu90, cu91, cu92, cu100, cu101, cu102, cu110–cu118, cu120–cu124, cu125, cu126, cu128, cu129, cu130, cu132	https://download.pytorch.org/whl/cuXXX
ROCm	rocm4.0.1, rocm4.1, rocm4.2, rocm5.1.1, rocm5.2–rocm5.7, rocm6.0–rocm6.4（含 rocm6.2.4）, rocm7.0, rocm7.1, rocm7.2	https://download.pytorch.org/whl/rocmX.Y
Intel XPU	xpu	https://download.pytorch.org/whl/xpu

所有索引 URL 均以 static LazyLock<IndexUrl> 形式硬编码在 backend.rs 尾部（例如 PYTORCH_CU128_INDEX_URL 指向 https://download.pytorch.org/whl/cu128），共 27 个 CUDA 版本 + 17 个 ROCm 版本 + CPU + XPU。

典型的实战用法：

# 自动探测本机驱动并选择索引
uv pip install torch --torch-backend auto

# 明确指定 CUDA 12.8 索引
uv add torch --torch-backend cu128

# 通过环境变量固定（适合 CI）
UV_TORCH_BACKEND=cpu uv sync

三、auto 模式的硬件探测链（accelerator.rs）

TorchMode::Auto 最终调用 Accelerator::detect，其文档注释列出了完整的探测顺序，优先级从高到低依次为：

1. UV_CUDA_DRIVER_VERSION 环境变量覆盖（对应 cuda_driver_version 参数）；
2. UV_AMD_GPU_ARCHITECTURE 环境变量覆盖（对应 amd_gpu_architecture 参数）；
3. 读取 /sys/module/nvidia/version——文件内容就是驱动版本，如 550.144.03；
4. 读取 /proc/driver/nvidia/version——从中提取第二段的驱动版本（格式示例见 parse_proc_driver_nvidia_version：NVRM version: NVIDIA UNIX Open Kernel Module for x86_64 550.144.03 Release Build ...）；
5. 执行 nvidia-smi --query-gpu=driver_version --format=csv,noheader，取首行作为驱动版本；
6. 执行 rocm_agent_enumerator，从输出行中解析出 AMD GPU 架构（如 gfx906）；
7. 遍历 /sys/bus/pci/devices，通过 PCI class 与 vendor 识别 Intel GPU（XPU）；
8. （仅 Windows）通过 Configuration Manager API 查询"当前在场"的 Intel 显示适配器设备树。

值得注意的细节：

• CUDA 驱动版本 ≠ CUDA 工具包版本。源码注释明确强调（accelerator.rs#L39-L42）：Cuda { driver_version } 存储的是驱动版本（如 550.144.03），而非 CUDA toolkit 版本（如 12.8.0）。这个区分是后续"驱动版本 → 可用 CUDA 后端"映射表的前提。
• 多 GPU 场景：nvidia-smi 在多个 GPU 上会输出多行版本号，代码只取 stdout.lines().next()（首行）。crate 内的单元测试 nvidia_smi_multi_gpu 专门验证了这一解析行为。
• 错误处理策略：探测过程中文件不存在（NotFound）会静默跳过进入下一手段，而其他 IO 错误或版本解析失败会向上返回 AcceleratorError；全部手段失败则返回 Ok(None)，最终回退到 CPU 索引。错误类型定义在 AcceleratorError（包含 Io、Version、Utf8、ParseInt、UnknownAmdGpuArchitecture 五类）。

3.1 AMD GPU 架构枚举

AmdGpuArchitecture 覆盖了 gfx900、gfx906、gfx908、gfx90a、gfx942、gfx950、gfx1030、gfx1100、gfx1101、gfx1102、gfx1150、gfx1151、gfx1200、gfx1201 共 14 种架构，并实现了 FromStr（"gfx906" → Gfx906，未知值报错）与 Display。多卡时取所有输出行中排序最小（.min()）的架构（accelerator.rs#L165-L170），即以最老的卡为准来保证兼容性。

3.2 Intel XPU 的 PCI 级检测

Linux 上通过读取每个 PCI 设备的 class 与 vendor 文件（parse_pci_device_ids）判断：PCI base class 为显示类（0x03，掩码 0x00ff_0000 比较）且 vendor 为 Intel（0x8086）即认定为 XPU。

Windows 路径（detect_intel_gpu_from_windows_devices）则调用 CM_Get_Device_ID_List_W 系列 API，以显示适配器类 GUID {4d36e968-e325-11ce-bfc1-08002be10318} 过滤"在场"设备，再检查设备实例 ID 中是否包含 VEN_8086 段（大小写不敏感，contains_intel_vendor_id）。该逻辑有对应的单元测试 intel_vendor_id_from_pnp_device_id 和 windows_device_instance_ids（accelerator.rs#L430-L459），验证了如 PCI\VEN_8086&DEV_9A49... 判为 Intel、PCI\VEN_10DE...（NVIDIA）判为非 Intel 等行为。

四、TorchStrategy：从 Mode 到策略的转换

TorchStrategy::from_mode 是核心转换函数：

pub fn from_mode(
    mode: TorchMode,
    os: &Os,
    cuda_driver_version: Option<Version>,
    amd_gpu_architecture: Option<AmdGpuArchitecture>,
) -> Result<Self, AcceleratorError>

• 若 mode 是 Auto：调用 Accelerator::detect，得到 Cuda { driver_version } / Amd { gpu_architecture } / Xpu 三种硬件策略之一；检测不到任何加速器则回退为 Backend { backend: Cpu }；
• 若 mode 是其余 51 个显式取值：直接一一映射为对应的 Backend { backend: ... }。

策略枚举 TorchStrategy 有四个变体：

变体	携带信息	语义
Cuda { os, driver_version }	操作系统 + 驱动版本	Linux/Windows 上按驱动版本筛选可用 CUDA 索引
Amd { os, gpu_architecture }	操作系统 + GPU 架构	按架构精确匹配 ROCm 索引
Xpu { os }	操作系统	仅 Manylinux/Windows 提供 XPU 索引，其余平台回退 CPU
Backend { backend }	显式后端	单一定死的索引

4.1 哪些包会被重定向到 PyTorch 索引

applies_to 列出了受 torch 策略管理的包名白名单，包括 torch、torchaudio、torchvision、torchtext、torchserve、torchrec、torchao、torchtune、torchcodec、torchcsprng、torchdistx、torch-tensorrt、torcharrow、triton、triton-rocm、triton-xpu、pytorch-triton、pytorch-triton-rocm、pytorch-triton-xpu、fbgemm-gpu、fbgemm-gpu-genai、xformers 等。只有这些包会被从用户配置的索引中"接管"。

4.2 哪些包还带有系统级依赖

has_system_dependency 是另一个更小的名单（torch、torchvision、vllm、flash-attn、flash-attn-3、deepspeed、megablocks、natten 等）。源码注释解释了区别：triton 虽然托管在 PyTorch 索引上，但它本身不依赖特定 CUDA 版本（cu128 索引上的 triton 并不绑定 CUDA 12.8），因此只有 applies_to 而没有 has_system_dependency。这个区分影响着 lockfile 中对系统依赖的标注与校验。

五、auto 模式下索引 URL 的最终生成逻辑

TorchStrategy::index_urls 实现了平台差异化的索引集合生成：

5.1 CUDA 策略（Self::Cuda）

• Manylinux / Musllinux：遍历静态表 LINUX_CUDA_DRIVERS（backend.rs#L767-L803），筛选出所有"最低驱动版本 ≤ 本机驱动版本"的 CUDA 后端，将其索引按从新到旧的顺序生成，最后再追加 CPU 索引作为兜底。例如本机驱动为 550.144.03，则可用的表项包括 cu120–cu129（要求驱动 ≥ 525.60.13）和 cu110–cu118（要求 ≥ 450.80.2），于是 uv 会依次尝试这些索引，并在解析失败时回退到 CPU。
• Windows：使用独立的 WINDOWS_CUDA_VERSIONS 表（backend.rs#L808-L844），因为 Windows 的 CUDA 驱动版本号体系与 Linux 不同（如 cu120–cu129 对应驱动 ≥ 528.33，cu110–cu118 对应 ≥ 452.39）。
• macOS / BSD / Android / Pyodide 等其他平台：直接只返回 CPU 索引——这些平台没有 CUDA 支持。

两张驱动映射表的注释指明其数据来源是 NVIDIA CUDA Toolkit 发布说明中的"Table 1/2"（CUDA 版本 ↔ 最低驱动版本），这与 light-the-torch 的做法一致。

5.2 AMD 策略（Self::Amd）

Manylinux/Musllinux 上遍历 LINUX_AMD_GPU_DRIVERS 表（backend.rs#L858-L964），该表按"ROCm 版本 × 受支持架构"逐条列出（ROCm 7.2 支持全部 14 种架构；ROCm 7.1/7.0 缺少 gfx1150/gfx1151；ROCm 6.4/6.3/6.2.4 缺少 gfx950；ROCm 6.2 及更早的列表则更短），筛选出架构完全匹配的 ROCm 后端，同样追加 CPU 索引兜底。表头的文档注释说明了数据获取方式：对每个 ROCm 版本运行 torch.cuda.get_arch_list() 实测推断，并交叉参考 AMD 官方兼容性矩阵。其他所有平台（含 Windows）则只返回 CPU 索引。

5.3 XPU 策略（Self::Xpu）

仅在 Manylinux 与 Windows 上返回 https://download.pytorch.org/whl/xpu；其余平台回退 CPU。

5.4 显式后端（Self::Backend）

无条件的单索引：backend.index_url()，即上表中该 backend 对应的固定 URL。

此外，TorchBackend::from_index 提供了反向解析：识别形如 https://download.pytorch.org/whl/cu124 的 URL（host 必须为 download.pytorch.org，且首段路径为 whl），将其第二段解析为 backend 字面量（cu124、rocm6.4、cpu、xpu 等，见 FromStr 实现）。这让 uv 能够识别用户在 pyproject.toml 或命令行中直接写死的 PyTorch 官方索引，并将其与 backend 体系对齐。cuda_version() / rocm_version() 两个方法则返回每个后端对应的 CUDA/ROCm 语义化版本，供依赖校验使用（backend.rs#L594-L703）。

六、版本兼容性与适用前提

• crate 版本对应关系：uv-torch 0.0.76 属于 uv 0.12.9 的组件（见 README 与 Cargo.toml）；作为内部 crate，其 API 不承诺稳定。
• preview 特性：--torch-backend 在 CLI 帮助中明确标注 "This option is in preview and may change in any future release"，即索引表、模式命名在未来版本中可能演进（例如新增 cu13x 或 ROCm 版本）。
• 探测依赖宿主环境：auto 模式依赖 /sys、/proc、nvidia-smi、rocm_agent_enumerator 等本机能力；在无 GPU 的 CI 容器上，探测会自然回退到 CPU 索引。若需要确定性行为，建议通过 --torch-backend 或 UV_TORCH_BACKEND 显式指定，或用 UV_CUDA_DRIVER_VERSION / UV_AMD_GPU_ARCHITECTURE 覆盖探测结果（这两个覆盖参数在 detect 中优先级最高）。
• CUDA 索引并非全量可用：TorchMode 中 CUDA 取值跳过了 cu127（源码枚举里从 cu126 直接到 cu128），这反映了 PyTorch 官方当前发布的版本组合，属于实现事实。

七、小结

uv-torch 用约 1500 行 Rust 代码完成了三件事：一套覆盖 NVIDIA/AMD/Intel 的多路径硬件探测（环境变量覆盖 → /sys//proc 文件 → nvidia-smi/rocm_agent_enumerator 外部命令 → PCI/Windows 设备树）；一套"驱动版本/架构 → 可用后端集合"的静态映射表（LINUX_CUDA_DRIVERS、WINDOWS_CUDA_VERSIONS、LINUX_AMD_GPU_DRIVERS）；以及"生态包白名单 + 系统依赖名单"两级过滤（applies_to / has_system_dependency）。对用户而言，只需记住三个入口：uv add torch --torch-backend auto（自动）、--torch-backend cu128（显式）、UV_TORCH_BACKEND（环境变量）；对贡献者而言，crates/uv-torch/src/accelerator.rs 与 backend.rs 中内嵌的单元测试（多 GPU 的 nvidia-smi 输出解析、PnP 设备 ID 判定、Windows 设备实例解码）为后续扩展新后端提供了清晰的测试范式。