Why Installing GPU Python Packages Is So Complicated

An ML package can install successfully and fail when its first GPU operation runs. The package build, host driver, and libraries needed by that operation must all be compatible.

Python wheels have no standard tag for GPU hardware requirements. Their Python version, ABI, and platform tags describe part of binary compatibility; they do not describe the GPU model, driver, or CUDA runtime a compiled extension needs.

PyTorch distributes CUDA, ROCm, and CPU builds with different compatibility requirements. Without GPU selection metadata, PyPI cannot choose the appropriate build for every deployment. Projects use indexes, package names, or extras to expose that choice.

Three different CUDA components matter:

Installing a toolkit does not install a working host driver. A prebuilt wheel can still compile some operations at first use; vLLM’s FlashInfer setup needs matching cached kernels or a toolkit for those operations.

The CUDA field in nvidia-smi describes the driver’s supported CUDA level, not an installed toolkit. A wheel’s CUDA runtime need not match that field exactly: NVIDIA’s compatibility rules determine which combinations run, with additional constraints for PTX compilation and newer features.

Check resolution, image builds, and GPU execution separately

Each stage answers a different question:

Choose the deployment backend explicitly during GPU-less image builds. Hardware detection on the build machine cannot identify the deployment GPU.

On the execution host, test an operation from each accelerator package the application needs. A PyTorch matrix multiplication does not exercise DeepSpeed’s compiled optimizer or vLLM’s inference kernels. Multi-GPU environments also need a communication check across processes.

How packages work around this today

Without standard metadata for GPU support, every CUDA-dependent project has invented its own distribution scheme, and a fourth scheme has grown up around the projects that never solved it.

Separate index URLs (the PyTorch approach)

PyTorch hosts its own package index at download.pytorch.org/whl/. Users select their CUDA version by pointing pip or uv at the matching URL:

# CUDA 12.8
pip install torch --index-url https://download.pytorch.org/whl/cu128

# CPU only
pip install torch --index-url https://download.pytorch.org/whl/cpu

The package name stays torch; the index selects a build family. Available Python versions, platforms, GPU architectures, and backends vary by release. Check the PyTorch selector rather than treating an index name as a promise that every wheel exists.

llama-cpp-python uses separate backend indexes too, with wheel selection depending on Python, CUDA, and GPU compute capability.

This approach keeps the package name clean but pushes complexity into project configuration. Every developer on a team needs to use the same index URL, and CI pipelines need separate configurations for GPU and CPU runners.

Package name suffixes (the RAPIDS approach)

NVIDIA’s RAPIDS libraries encode the CUDA family in package names such as cudf-cu12, cuml-cu12, and cudf-cu13. Current releases publish wheels on PyPI; NVIDIA’s pypi.nvidia.com index is also used by some releases. NVIDIA’s lower-level runtime libraries use the same convention: nvidia-cuda-runtime-cu12.

pip install cudf-cu12

Users must know their CUDA major version and select the matching package. Dependency resolution becomes awkward because cudf-cu12 and cudf-cu13 look like unrelated packages to the installer, even though they provide the same Python API.

Extras syntax (the JAX approach)

JAX uses pip extras to pull in the right backend:

pip install "jax[cuda13]"

The extras trigger installation of CUDA-specific dependencies. This keeps the base package name intact and uses a standard pip mechanism, though it still requires the user to specify the CUDA version manually.

Curated rebuild indexes (the Astral approach)

The first three schemes are things a publisher does to its own releases. The fourth is something a third party does on the publisher’s behalf.

Packages like flash-attn and DeepSpeed compound the matching problem: a working binary has to line up with the Python version, the PyTorch version, the CUDA version, and the C++ ABI at once. flash-attn’s setup.py responds by guessing a wheel filename from the installed environment, then falling back to a multi-hour compile when the guess misses.

The Astral GPU indexes rebuild those packages and move both moving parts into a PEP 440 local version segment, with one index per CUDA version:

flash_attn-2.8.3.post1+cu.12.8.torch.2.8-cp312-cp312-manylinux_2_24_x86_64.whl

The local version identifies the CUDA and PyTorch build pair. The inspected FlashAttention wheels also constrain the PyTorch minor version through a dependency such as torch==2.11.*. Add flash-attn to [project].dependencies, pin a matching wheel, and route it to the curated index:

[tool.uv.sources]
flash-attn = { index = "astral-cu128" }

[[tool.uv.index]]
name = "astral-cu128"
url = "https://wheels.astral.sh/simple/cu128/"
explicit = true

The manifest lists supported packages and combinations, including flash-attn, DeepSpeed, vLLM, and torch-scatter. Coverage differs across CUDA, Python, architecture, and PyTorch versions; check the actual wheel rather than assuming every combination is available.

Three limits keep this a patch rather than a fix.

It covers a curated set of packages. Projects outside that set still need their own binary distribution or source builds.

It removes some source builds, but runtime behavior still matters. Packages such as DeepSpeed probe optional ops or JIT-compile extensions. A prebuilt wheel does not establish that every operation works without a toolkit.

It does not fully enforce the CUDA half of the pair. A dependency such as torch==2.11.* accepts different local CUDA builds of that PyTorch release. The extension’s +cu.12.8 label does not force torch to come from the CUDA 12.8 index. Route PyTorch to a matching index too, then verify imports and an accelerator operation.

That last limit is the argument for solving this in the format. A curated index needs a maintainer and covers whatever that maintainer chooses to build, while wheel variants put the same selection into metadata any publisher can emit.

How uv handles this today

uv has built-in support for routing packages to different indexes based on platform markers. Add torch to [project].dependencies before adding this configuration to pyproject.toml: source mappings apply to direct dependencies, so a transitive PyTorch dependency alone does not use this routing.

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
]

This routes CUDA builds to Linux and CPU builds elsewhere. Operating-system markers do not detect a GPU: Linux CPU hosts still receive CUDA wheels. uv’s --torch-backend=auto queries visible drivers and accelerator information in the uv pip interface, falling back to CPU when no accelerator is detected.

For a GPU-less image build, select the deployment backend explicitly. Building and importing packages without a GPU checks packaging; execution still needs a GPU test that asserts device availability and completes a small operation. Auto detection during a CPU-only build cannot predict the deployment hardware.

What conda and pixi do differently

Conda sidesteps this problem by managing CUDA as a non-Python dependency. A conda environment file can pin a specific CUDA toolkit version alongside pytorch and numpy, and conda’s solver resolves all of them together. The CUDA toolkit is just another package in the dependency graph.

Pixi inherits this capability through conda-forge. The solver coordinates CUDA package versions and host capability requirements, but it does not install a host driver or prove that a GPU operation works. Wheels added alongside conda packages still need compatible binaries; manage the boundary between conda and pip packages explicitly.

What wheel variants will change

PEP 817, the wheel variant model, describes providers and selection beyond platform tags. PEP 825, the variant package format, specifies how wheels carry variant metadata. PEP 817 is a draft informational proposal; PEP 825 is provisionally accepted, with final acceptance depending on the remaining standards.

The key idea: wheels gain structured metadata like nvidia :: cuda_version_lower_bound :: 12.8. At install time, a provider plugin queries the local hardware (GPU model, driver version) and reports what variants the machine supports. The installer then picks the best match from available wheels.

Variants are not a GPU-only mechanism. The same property system describes CPU microarchitecture levels (x86_64 :: level :: v3 covers the AVX2-and-FMA baseline) and other compatibility dimensions the current platform tags cannot express. For the full explanation of how the format and plugin model work, see What are wheel variants?.

With publisher and installer support, pip install torch or uv add torch could select an appropriate CUDA build without manual index configuration. Astral’s experimental variant-enabled build of uv demonstrates this selection; it is separate from the mainline workflow documented here.

What to do now

For PyTorch-centric projects, configure index URLs in pyproject.toml using uv’s source routing. This gives reproducible builds without manual index URL flags. For the packages that compile against PyTorch, check the Astral GPU indexes before reaching for a source build. For mixed GPU workloads spanning PyTorch and RAPIDS, consider whether conda or pixi would reduce the configuration burden, since both manage CUDA versions as first-class dependencies.

Follow the wheel variant standards and installer implementations. Format acceptance alone does not replace index routing: publishers must supply variant wheels, and installers must support selecting them.