Skip to content

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

Open
golithe wants to merge 2 commits into
replicate:mainfrom
golithe:gpu-compat-doctor-check
Open

feat: detect when the resolved torch/CUDA ships no kernels for the local GPU#3104
golithe wants to merge 2 commits into
replicate:mainfrom
golithe:gpu-compat-doctor-check

Conversation

@golithe

@golithe golithe commented Jul 14, 2026

Copy link
Copy Markdown
  • 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 (lowest across GPUs, since the image must run on the weakest) 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: 2.7.0+cu118 is a genuine 2.7 build with no Blackwell kernels
  • 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

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

@golithe
golithe requested a review from a team as a code owner July 14, 2026 08:31
@golithe golithe changed the title feat: detect when the resolved torch/CUDA ships no kernels for the lcal GPU feat: detect when the resolved torch/CUDA ships no kernels for the local GPU Jul 14, 2026

@anish-sahoo anish-sahoo left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread pkg/doctor/check_env_gpu_compat.go Outdated
Comment thread pkg/doctor/check_env_gpu_compat.go Outdated
Comment thread pkg/doctor/check_env_gpu_compat.go Outdated
Comment thread pkg/doctor/check_env_gpu_compat_test.go Outdated
Comment thread pkg/doctor/gpu_compatibility.go
Comment thread pkg/doctor/gpu_compatibility.go
@golithe
golithe requested a review from anish-sahoo July 17, 2026 13:09
@golithe

golithe commented Jul 17, 2026

Copy link
Copy Markdown
Author

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

@anish-sahoo anish-sahoo left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants