feat: detect when the resolved torch/CUDA ships no kernels for the local GPU by golithe · Pull Request #3104 · replicate/cog

golithe · 2026-07-14T08:31:31Z

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

…cal GPU

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.

anish-sahoo · 2026-07-16T17:29:08Z

+		}
+		return caps[i][1] < caps[j][1]
+	})
+	return caps[0], true


Could we evaluate every detected capability, or select the capability with the strictest applicable floor? Sorting ascending and returning only caps[0] checks the numerically lowest GPU, which is not necessarily the least compatible one. On a host with sm_90 and sm_120, torch 2.4.1/CUDA 12.4 passes the sm_90 floor and produces no finding even though operations on sm_120 will fail. A mixed-GPU regression test would help cover this.

the check now evaluates every distinct capability rather than only the lowest. your described scenario (sm_90 + sm_120 with torch 2.4.1 is the regression test TestGPUCompatibilityCheck_MixedGPUsFireForStrictestFloor

wdyt?

anish-sahoo · 2026-07-16T17:29:08Z

+		return nil, nil
+	}
+
+	torchVersion, hasTorch := ctx.Config.TorchVersion()


TorchVersion() does not establish that this is an exact pin: the current requirement parsing discards the comparator and returns the first version token. As a result, torch<2.7.0 is treated as 2.7.0, so this check can pass even though pip installs an older, incompatible release. Could we evaluate only normalized exact == pins and skip ranges, direct URLs, and unresolved requirements? This should also have regression coverage.

added requirements.ExactVersion and Config.TorchExactVersion where only a single == specifier counts

ranges/wildcards/direct URLs/unpinned all produce no finding rather than a guess

tested in TestExactVersion and TestGPUCompatibilityCheck_SkipsNonExactPin

anish-sahoo · 2026-07-16T17:29:08Z

+		return nil
+	}
+
+	torchOK := version.GreaterOrEqual(torchVersion, floor.MinTorch)


version.GreaterOrEqual includes local metadata in equality, so 2.7.0+cu128 is not considered greater than or equal to 2.7.0: the release is not greater, and equality fails because the metadata differs. That produces a false warning at the exact supported floor. Could we strip the local modifier for the torch release comparison, retain it separately for CUDA resolution, and add an exact-boundary test?

the torch comparison now strips the local tag before version.GreaterOrEqual, so 2.7.0+cu128 passes at the exact floor

CUDA bound is still checked separtely from Build.CUDA

boundari case in TestEvaluateGPUCompat ("sm_120 exact floor with local tag passes)

anish-sahoo · 2026-07-16T17:29:08Z

+func TestGPUCompatibilityCheck_RunsWithoutError(t *testing.T) {
+	// A GPU may or may not be present in the test environment; just ensure no panic or error.
+	ctx := gpuContext(t, true, "2.4.1", "12.4")
+	_, err := (&GPUCompatibilityCheck{}).Check(ctx)


This test only verifies that the check returns no error and ignores its findings, so the production detection path is not exercised deterministically. Could we test it with a fake nvidia-smi and cover multiple GPUs, malformed output, command failure, and registration through the doctor runner? That would also catch the mixed-GPU issue noted above.

replaced the RunsWithoutError test entirely

the nvidia-smi probe is now a package-level functin variable that tests stub

went with that over fake bin on PATH to keep it hermetic, same intent

now covers multiple GPUs, command failure, malformed output, non-exact pins, and runner registration

anish-sahoo · 2026-07-16T17:29:08Z

+func (c *GPUCompatibilityCheck) Description() string { return "GPU compatibility" }
+
+func (c *GPUCompatibilityCheck) Check(ctx *CheckContext) ([]Finding, error) {
+	if os.Getenv("COG_SKIP_GPU_CHECK") != "" {


Could we document COG_SKIP_GPU_CHECK in docs/environment.md alongside COG_SKIP_DOCKER_CHECK, then regenerate docs/llms.txt? This is a new user-facing environment variable and its intended use when building for different hardware should be discoverable outside the source.

anish-sahoo · 2026-07-16T17:29:08Z

@@ -0,0 +1,199 @@
+package doctor


Could we rename these files to gpu_compatibility.go and gpu_compatibility_test.go? We prefer subject-based filenames here and spell out compatibility rather than shortening it to compat.

golithe · 2026-07-17T13:10:29Z

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

feat: detect when the resolved torch/CUDA ships no kernels for the lo…

70e49af

…cal GPU

golithe requested a review from a team as a code owner July 14, 2026 08:31

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 requested changes Jul 16, 2026

View reviewed changes

Address reviewer comments

3a7e9f1

golithe requested a review from anish-sahoo July 17, 2026 13:09

Uh oh!

Conversation

golithe commented Jul 14, 2026

Uh oh!

anish-sahoo left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

golithe commented Jul 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants