fix: TensorRT-RTX Turing (SM 7.5) capability guards and compute-capability targeting - #4643
Open
tp5uiuc wants to merge 3 commits into
Open
fix: TensorRT-RTX Turing (SM 7.5) capability guards and compute-capability targeting#4643tp5uiuc wants to merge 3 commits into
tp5uiuc wants to merge 3 commits into
Conversation
This was referenced Aug 29, 2026
tp5uiuc
force-pushed
the
tp5uiuc/trtrtx-turing-foundation
branch
2 times, most recently
from
August 29, 2026 23:01
0828d80 to
60522c1
Compare
tp5uiuc
commented
Aug 30, 2026
tp5uiuc
force-pushed
the
tp5uiuc/trtrtx-turing-foundation
branch
from
August 30, 2026 00:49
60522c1 to
bb9c235
Compare
Adds a target_compute_capabilities setting naming the architectures an engine is built for, and declares them on the TensorRT-RTX builder config, so an artifact can be produced for something other than the build host. Declaring it on standard TensorRT raises: that backend always builds for the current device. An undeclared target is now resolved rather than skipped. Left at the builder default num_compute_capabilities stays 0, and a refittable graph then fails with "Compatible cubin or ptx module for device target '75' not found" -- so on Turing every refittable build failed. It now resolves to the current device the way partitioning already did via get_target_compute_capabilities(), naming ComputeCapability.CURRENT rather than an SM<major><minor> lookup because TensorRT-RTX names only a subset of architectures. The builder-config work lives in set_rtx_compute_capabilities() so it can be unit tested against a stub config, with no GPU and no engine build. Not exposed on cross_compile_for_windows: TensorRT-RTX does not support that path. The setting is engine-invariant: an engine built for one capability set must not be reused for a compile targeting another.
TensorRT-RTX does not support FP32 GEMMs or 3D convolutions on Turing (SM 7.5), and Turing has no bfloat16 hardware. Handed those ops anyway it returns a null execution context, segfaults on bf16, or -- under dynamic shapes -- builds, runs and returns an all-zero tensor with no exception. That silent case is the motivating one. Guards key off the capabilities being built for rather than the build host, so an ahead-of-time build for another architecture partitions correctly. The GEMM guard keys on fp32 operands only; the convolution guard covers forward 3D only, since transposed 3D works on Turing. bfloat16 is gated in the partitioners because the crash is not operator-specific, mirroring the existing complex-dtype handling. Converter unit tests need explicit skips: DispatchTestCase bypasses the partitioner, so a guarded node raises UnsupportedOperatorException instead of falling back and an unguarded one reaches TensorRT-RTX and fails. The cdist skips state the converter's condition rather than p == 2, because a GEMM is only emitted for compute_mode 1, or 0/absent with an operand above the row threshold. Measured on a T4: the 3 cases outside that condition pass, the 9 inside it fail. Rather than hand-copy the branch, it is extracted from cdist_forward as cdist_emits_matmul/CDIST_MATMUL_ROW_THRESHOLD -- the converter now consumes its own predicate, and the test and the Turing capability validator consume the same one, so none of the three can drift. The threshold was a bare literal in the converter until now. Four parameterisations named for a compute_mode they do not use are renamed. Also clears two pre-existing lint failures in aten_ops_converters.py that pre-commit blocks on now the file is in the changed set: an unused type: ignore[assignment] and a "mis-evaluate" spelling.
Adds the regression tests for the fallbacks introduced alongside the Turing (SM 7.5) capability guards: FP32 GEMM, 3D convolution and bfloat16 must fall back to PyTorch, while FP16 GEMM, 2D convolution and transposed 3D convolution must stay on TensorRT. Most of the coverage is written against target_compute_capabilities=[(7, 5)], which forces Turing's partitioning on any GPU, so the guards are exercised in CI without Turing hardware. A second class repeats the same checks natively and is skipped off SM 7.5; it also pins the two failure modes that motivated the guards -- a null execution context under static shapes and a silently all-zero result under dynamic shapes. Two further cases assert on the builder config rather than on compile success, which is what makes them meaningful off Turing: everywhere except SM 7.5 an undeclared compute capability is silent.
tp5uiuc
force-pushed
the
tp5uiuc/trtrtx-turing-foundation
branch
from
August 30, 2026 18:37
bb9c235 to
633444a
Compare
Collaborator
Author
CI summary
Suites ran to completion: Why these are not from this PR: the same |
This was referenced Aug 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What — Adds a
target_compute_capabilitiescompilation setting, declares it on the TensorRT-RTXbuilder config, and falls back to PyTorch for the ops TensorRT-RTX cannot serve on Turing.
Why — TensorRT-RTX runs on SM 7.5 and up, but its support matrix carves Turing out of several
paths: FP32 GEMMs and 3D convolutions are unsupported at compute capability 7.5, and Turing has no
bfloat16 hardware. Torch-TensorRT had no notion of this and handed those ops over anyway:
createExecutionContext()returns nullcorrect shape and dtype, with no exception
The dynamic-shape GEMM is the motivating case, because it fails silently.
Separately, the builder was never told which architecture it was building for. With
num_compute_capabilities == 0, Myelin looks for a precompiled module instead of JIT-ing one andthe build dies —
Compatible cubin or ptx module for device target '75' not found— so on Turingevery refittable build failed.
How — Guards key off the capabilities being built for, not the build host, so an
ahead-of-time build for another architecture partitions correctly instead of baking in the build
machine. An undeclared target is resolved rather than skipped, to the current device, matching how
partitioning already resolved it — and by naming
ComputeCapability.CURRENTrather than looking upSM<major><minor>, because TensorRT-RTX names only a subset of architectures and the implicit pathmust not start failing on hosts that build fine today.
The GEMM guard keys on fp32 operands only, so fp16 GEMMs accumulating in fp32 are unaffected. The
convolution guard covers forward 3D only, since transposed 3D works on Turing. bfloat16 is gated at
the partitioner rather than per-converter because the crash is not operator-specific, mirroring the
existing complex-dtype handling. Converter unit tests bypass the partitioner, so a rejected node
raises instead of falling back; those tests skip explicitly.
Testing — Confirmation sweep at the series tip, T4 (SM 7.5) and L40S (SM 8.9), driver 595.58.03,
identical stacks: T4 2668 passed / 17 failed / 226 skipped; L40S 2780 / 13 / 118; 2911 collected on
each. Declaring the capability alone closes 42 of the 91 original Turing failures. Defaulting
costs nothing off Turing: on an L40S, engines built with the default and with an explicit
[(8, 9)]are byte-identical (280148 B and 287100 B for a small conv net; 47148116 B and 24184052 B for
resnet18), with no measurable build-time difference. On the L40S there is not one status change
across all six test modules — every guard is inert off Turing.
Cost / Gotchas
conversion/: the T4 runs 1981 of2126 tests where the L40S runs 2079. Green on Turing does not mean equally tested on Turing.
aten_ops_converters.py(an unusedtype: ignore[assignment]and a spelling thetyposhook rejects). Both predate this change, butthe file is now in the changed set so pre-commit blocks on them.
Followups — Three more capability gaps ship as stacked PRs on this branch: 3D convolution behind
the pad-folding pass, the GEMM
cdistemits internally, and FP32 GEMMs reaching TensorRT-RTX throughlinearand attention.🤖 Generated with Claude Code