From 3e34aee63efb8b20eb035ba32706a18e7babae00 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 22 Aug 2025 12:06:23 +0100
Subject: [PATCH] Avoid introducing unnecessary system dependency on CUDA
 (#15449)

## Summary

Packages like `triton` should come from the PyTorch index, but they
don't actually vary across (e.g.) the `cu128` or `cu129` indexes.

Closes https://github.com/astral-sh/uv/issues/15446.

## Test Plan

Validate that the following pins to `cu128`, rather than `cpu`:

```
echo "vllm\ntorch==2.7.1+cu128" | cargo run pip compile --torch-backend=auto --extra-index-url https://wheels.vllm.ai/b2f6c247a9b84556a8ea0e75bb4a2db765ff3315 - --python-platform linux --python-version 3.13 -v
```
---
 crates/uv-resolver/src/resolver/mod.rs |  1 +
 crates/uv-torch/src/backend.rs         | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
index d0f43e9c9..92f9fba38 100644
--- a/crates/uv-resolver/src/resolver/mod.rs
+++ b/crates/uv-resolver/src/resolver/mod.rs
@@ -1839,6 +1839,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     .torch_backend
                     .as_ref()
                     .filter(|torch_backend| matches!(torch_backend, TorchStrategy::Cuda { .. }))
+                    .filter(|torch_backend| torch_backend.has_system_dependency(name))
                     .and_then(|_| pins.get(name, version).and_then(ResolvedDist::index))
                     .map(IndexUrl::url)
                     .and_then(SystemDependency::from_index)
diff --git a/crates/uv-torch/src/backend.rs b/crates/uv-torch/src/backend.rs
index 6f050e3a5..12e69c181 100644
--- a/crates/uv-torch/src/backend.rs
+++ b/crates/uv-torch/src/backend.rs
@@ -277,6 +277,25 @@ impl TorchStrategy {
         )
     }
 
+    /// Returns `true` if the given [`PackageName`] has a system dependency (e.g., CUDA or ROCm).
+    ///
+    /// For example, `triton` is hosted on the PyTorch indexes, but does not have a system
+    /// dependency on the associated CUDA version (i.e., the `triton` on the `cu128` index doesn't
+    /// depend on CUDA 12.8).
+    pub fn has_system_dependency(&self, package_name: &PackageName) -> bool {
+        matches!(
+            package_name.as_str(),
+            "torch"
+                | "torcharrow"
+                | "torchaudio"
+                | "torchcsprng"
+                | "torchdata"
+                | "torchdistx"
+                | "torchtext"
+                | "torchvision"
+        )
+    }
+
     /// Return the appropriate index URLs for the given [`TorchStrategy`].
     pub fn index_urls(&self) -> impl Iterator<Item = &IndexUrl> {
         match self {