feat: detect when the resolved torch/CUDA ships no kernels for the local GPU

[golithe]

• Cog resolves CUDA from the framework pin and never consults the GPU it runs on
• On a Blackwell card (sm_120), torch==2.4.1 builds green, and every kernel launch fails with CUDA error: no kernel image is available for execution on the device
• Hard to notice: torch.cuda.is_available() returns True, and PyTorch's own diagnostic is a UserWarning
• There's no fallback either, since the wheel ships neither sm_12x cubins nor PTX to JIT from
• cog run does not work with A100 #389 is the same failure class on an A100 in 2022 (closed when the reporter worked around it); Blackwell makes it current again

This PR adds:

• GPUCompatibilityCheck to cog doctor, modelled on DockerCheck
• Reads compute capability from nvidia-smi --query-gpu=compute_cap (every distinct GPU, not just the weakest, since a mixed-GPU host must clear each device's floor) and compares the resolved torch/CUDA against the oldest release known to ship kernels for it.

cog doctor output when it fires:

 ⚙  Environment
 ✔ Docker
 ✔ Python version
 ⚠ GPU compatibility
 ⚙    torch==2.4.1 (CUDA 12.4) ships no kernels for sm_120, the compute capability of
      this machine's GPU. The image will build, but every CUDA operation in it will
      fail at runtime with "no kernel image is available for execution on the device".

 ⚙  Found 1 warning.

Notes for review:

• The floors are measured, not read off release notes. I installed each wheel and read torch._C._cuda_getArchFlags(), which works without a GPU
• Each floor is bracketed: the named version ships the kernels, the release below it doesn't. sm_120/sm_100 need torch>=2.7.0 + CUDA>=12.8; sm_90 needs torch>=2.0.1 + CUDA>=11.8
• Both bounds are load-bearing: a 2.7 build against an old CUDA (cu118) has no Blackwell kernels. The CUDA compared is the one from the wheel Cog actually resolves and installs (same resolution as PythonRequirementsForArch), not the raw requirement, so it's correct in both directions: torch==2.7.0+cu118 with cuda: "12.8" resolves to the cu128 wheel and stays silent, while an explicit --extra-index-url=.../cu118 installs the cu118 wheel regardless of cuda: and warns against 11.8
• Deliberately no rows below sm_90. Every probed wheel already covers Turing/Ampere/Ada, so those floors can't be bracketed; unknown-old devices produce no finding rather than a guess
• SeverityWarning, matching PythonVersionCheck: the image is valid and runs fine on other hardware; it only fails when executed on this machine's GPU
• Silent when there's no GPU, no torch pin, or gpu: false; COG_SKIP_GPU_CHECK=1 skips it when building for different hardware than the local card. Doctor-only, no build-path changes
• The torch pin is matched PEP 503-style, so Torch== and torch[extra]== resolve rather than silently skipping the check.

Testing:

• the comparison is a pure function (evaluateGPUCompat), table-tested without a GPU, including +cu128 local-tag pins and the new-torch/old-CUDA case
• Verified by hand on an RTX 5070 Ti (sm_120): fires on torch==2.4.1, silent on 2.7.1 and on gpu: false

Out of scope (deliberately):

• the floors are hand-entered from measurement (precedent: MinimumTorchVersion and friends in pkg/dockerfile/base.go)
• A probe script in tools/ to regenerate them on demand could be a follow-up; full compatgen integration seems like a poor fit, since new capability majors ship every couple of years and the probe needs multi-GB wheel installs. Happy to add the script here or in a follow-up PR if useful

[anish-sahoo]

Thanks for putting this together. This is a useful diagnostic and the overall approach looks promising. I found several compatibility edge cases, along with some test and documentation gaps, that need additional fixes before we can work toward merging this. Details are inline.

[golithe]

@anish-sahoo ty for your thorough review! i addressed your comments, PTAL

[anish-sahoo]

Thanks for addressing the first review round. The mixed-GPU evaluation, exact-pin handling for ordinary torch==... requirements, local-tag boundary case, deterministic probe tests, docs, and rename all look good to me.

I found two remaining issues that need to be addressed before merge:

1. gpu_compatibility.go strips the local wheel tag and then relies only on build.cuda. This can silence a real Blackwell incompatibility: torch==2.7.0+cu118 with an explicit cuda: "12.8" passes the check because it compares the stripped 2.7.0 and base CUDA 12.8 against the floor. But Cog retains the +cu118 wheel, and that wheel has no Blackwell kernels regardless of its CUDA base image. Please preserve/derive the wheel CUDA tag for compatibility evaluation, rather than losing it when normalizing the torch release, and add a full Config.Complete -> Check regression test for this override combination.

2. Config.TorchExactVersion() compares requirements.PackageName(pkg) directly to lowercase "torch". Exact valid requirements such as torch[extra]==2.4.1 and Torch==2.4.1 therefore return no torch pin and skip the GPU check. Please normalize the base distribution name (strip extras and normalize casing, ideally PEP 503-style) before matching, with an end-to-end check regression test.

[golithe]

gpu_compatibility.go strips the local wheel tag and then relies only on build.cuda.

good catch. now derives the wheel's CUDA from the +cuXXX local tag (via the exported config.CUDAVersionFromTorchLocalTag) and compares that against the floor instead of build.cuda, so torch==2.7.0+cu118 with cuda: "12.8" fires. regression: TestGPUCompatibilityCheck_LocalCUDATagOverridesBuildCUDA

Config.TorchExactVersion() compares requirements.PackageName(pkg) directly to lowercase "torch".

fixed. added requirements.NormalizePackageName (PEP 503: strip extras, collapse -_., lowercase) and match on it in pythonPackageExactVersion, so torch[extra]==2.4.1 and Torch==2.4.1 resolve instead of being skipped. Regression: TestNormalizePackageName + end-to-end TestGPUCompatibilityCheck_NormalizesTorchName.

[golithe]

@anish-sahoo adressed both, PTAL

[anish-sahoo]

The GPU compatibility check can evaluate a different wheel than Cog actually installs, which causes both false positives and false negatives. Everything else LGTM.

[anish-sahoo]

I re-tested this and need to correct my earlier assumption. The GPU check and the build resolve CUDA differently:

• torch==2.7.0+cu118 with cuda: "12.8" builds with the cu128 wheel, but the check uses cu118 and warns incorrectly.
• An explicit cu118 index builds with the cu118 wheel, but the check uses cuda: "12.8" and does not warn.

Could the check use the final resolved requirement and index from PythonRequirementsForArch? That would ensure it checks the wheel Cog will actually install.

[golithe]

good catch again, and you're right.

• the check now resolves the wheel via the same path as PythonRequirementsForArch (new Config.ResolvedTorchWheel) and derives CUDA from the resolved package index instead of the requirement's local tag, so it matches what Cog installs in both directions: torch==2.7.0+cu118 + cuda: "12.8" resolves to cu128 and stays silent, and an explicit --extra-index-url=.../cu118 fires on 11.8.
• Regressions for both in TestResolvedTorchWheel and the two GPUCompatibilityCheck tests. PTAL.