fix: TensorRT-RTX Turing (SM 7.5) capability guards and compute-capability targeting by tp5uiuc · Pull Request #4643 · pytorch/TensorRT

tp5uiuc · 2026-08-29T17:53:17Z

What — Adds a target_compute_capabilities compilation setting, declares it on the TensorRT-RTX
builder 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:

FP32 GEMM, static shapes — createExecutionContext() returns null
FP32 GEMM, dynamic shapes — the engine builds and runs, returning an all-zero tensor of the
correct shape and dtype, with no exception
3D convolution — null execution context
bfloat16 — segmentation fault

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 and
the build dies — Compatible cubin or ptx module for device target '75' not found — so on Turing
every 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.CURRENT rather than looking up
SM<major><minor>, because TensorRT-RTX names only a subset of architectures and the implicit path
must 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

The guards switch off 105 tests on Turing, 97 of them in conversion/: the T4 runs 1981 of
2126 tests where the L40S runs 2079. Green on Turing does not mean equally tested on Turing.
Also clears two pre-existing lint failures in aten_ops_converters.py (an unused
type: ignore[assignment] and a spelling the typos hook rejects). Both predate this change, but
the 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 cdist emits internally, and FP32 GEMMs reaching TensorRT-RTX through
linear and attention.

🤖 Generated with Claude Code

tp5uiuc · 2026-08-29T17:55:04Z

Stack — merge #4643 first; #4644, #4645 and #4646 are based on it and GitHub will retarget them to main once it lands. They are independent of each other and can merge in any order, but all three add cases to tests/py/dynamo/models/test_turing_capability_guards.py, so the second and third will want a trivial rebase. #4647 and #4648 are not Turing bugs and are independent of this stack entirely.

tp5uiuc · 2026-08-30T00:03:42Z

+ cannot see shapes fail open. Shared with the capability validator that rejects this
+ case on Turing and with the converter tests, so none of them can drift from the
+ branch below.


Remove "Shared with the capability validator that rejects this
case on Turing and with the converter tests, so none of them can drift from the
branch below."

tp5uiuc · 2026-08-30T00:04:59Z

+# host, so a module compiled on Ampere and shipped to Turing does not retain ops Turing
+# cannot execute.
+
+TURING_COMPUTE_CAPABILITY = (7, 5)


Push it to inside the function, no one else consumes it.

tp5uiuc · 2026-08-30T00:07:24Z

+ getattr(settings, "target_compute_capabilities", None) if settings else None
+ )
+ if targets:
+ return tuple((int(major), int(minor)) for major, minor in targets)


tuple(tuple(map(int, major_minor)) for major_minor in targets)

tp5uiuc · 2026-08-30T00:09:12Z

+ targets = (
+ getattr(settings, "target_compute_capabilities", None) if settings else None
+ )
+ if targets:


if settings and (targets := getattr(settings, "target_compute_capabilities", None)):

tp5uiuc · 2026-08-30T00:11:18Z

+ # Declared targets name the architectures the artifact is built for; Turing is
+ # opt-in as an additional target because it can constrain kernel selection for the
+ # others. An undeclared target is resolved rather than skipped: the builder default
+ # leaves num_compute_capabilities at 0 and a refittable graph then fails with
+ # "Compatible cubin or ptx module for device target '75' not found". Resolve it to
+ # the current device the way get_target_compute_capabilities() already does for the
+ # validators, naming ComputeCapability.CURRENT because TensorRT-RTX names only a
+ # subset of architectures.
+ if ENABLED_FEATURES.tensorrt_rtx:
+ declared = self.compilation_settings.target_compute_capabilities
+ if declared:
+ builder_config.num_compute_capabilities = len(declared)
+ for idx, (major, minor) in enumerate(declared):
+ name = f"SM{major}{minor}"
+ compute_capability = getattr(trt.ComputeCapability, name, None)
+ if compute_capability is None:
+ supported = [
+ m for m in dir(trt.ComputeCapability) if m.startswith("SM")
+ ]
+ raise ValueError(
+ f"TensorRT-RTX has no compute capability {name} for requested "
+ f"target ({major}, {minor}). Supported: {supported}"
+ )
+ if not builder_config.set_compute_capability(
+ compute_capability, idx
+ ):
+ raise RuntimeError(
+ f"Failed to set TensorRT-RTX compute capability {name}"
+ )
+ _LOGGER.info(f"Targeting TensorRT-RTX compute capabilities {declared}")
+ else:
+ builder_config.num_compute_capabilities = 1
+ if not builder_config.set_compute_capability(
+ trt.ComputeCapability.CURRENT, 0
+ ):
+ raise RuntimeError(
+ "Failed to set the TensorRT-RTX compute capability of the "
+ "current device"
+ )
+ _LOGGER.info(
+ "Targeting the TensorRT-RTX compute capability of the current "
+ f"device {get_target_compute_capabilities(self.compilation_settings)}"
+ )


This should be refactored into its own function and unit-tested if possible.

Also the if else path (between declared SM and current) can also be two different hidden functions for ease of readability. so the top level call can be
if ENABLED_FEATURES.tensorrt_rtx:
set_compute_capabilities(self.build_config, self.compilation_settings.target_compute_capabilities)
and the rest can cascade into own function calls

tp5uiuc · 2026-08-30T00:30:35Z

 enable_experimental_decompositions (bool): Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT.
 dryrun (bool): Toggle for "Dryrun" mode, running everything except conversion to TRT and logging outputs
 hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer)
+ target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. Drives both engine targeting and op partitioning, so ops unsupported on any listed target fall back to PyTorch.


Cab we add a check where we error out if target_compute_capabilities is set for standard TensorRT (saying it is not supported)

This support check can be here or in COmpilationSettings post init, wherever it makes most sense

tp5uiuc · 2026-08-30T00:30:47Z

+# None means "target the current device". Set explicitly to build an engine deployable
+# on other architectures; see torch_tensorrt._utils for why capability validators must
+# consult this rather than the build host's device.


tp5uiuc · 2026-08-30T00:31:02Z

+ # Bind to a typed local: the imported helper is untyped from mypy's view here,
+ # and returning it directly trips --strict's no-any-return.


tp5uiuc · 2026-08-30T00:31:13Z

- return validate_disabled_constant_fold_exclusions(rule_ids)
+ # Bind to a typed local: the imported helper is untyped from mypy's view here,
+ # and returning it directly trips --strict's no-any-return.
+ normalized: Set[str] = validate_disabled_constant_fold_exclusions(rule_ids)


Note : this is a mypy fix.

tp5uiuc · 2026-08-30T00:33:20Z

+ # An engine built for one set of compute capabilities cannot be reused for a
+ # compile targeting a different set -- doing so would silently reintroduce ops the


Reemove this comment

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.

meta-claBot added the cla signed label Aug 29, 2026

github-actionsBot requested a review from zewenli98 August 29, 2026 17:53

tp5uiuc mentioned this pull request Aug 29, 2026
fix: fall back to PyTorch for ops TensorRT-RTX cannot serve on Turing (SM 7.5) #4546
Closed
7 tasks

This was referenced Aug 29, 2026
fix: guard the 3D convolutions the pad-folding pass hides from the conv guard #4644
Open
fix: guard the GEMM the cdist converter emits on Turing #4645
Open
fix: guard the FP32 GEMMs that reach TensorRT-RTX through linear and attention #4646
Open

tp5uiuc self-assigned this Aug 29, 2026

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-foundation branch 2 times, most recently from 0828d80 to 60522c1Compare August 29, 2026 23:01

github-actionsBot added the component: converters Issues re: Specific op converters label Aug 29, 2026

tp5uiuc commented Aug 30, 2026
View reviewed changes

tp5uiuc added 3 commits August 29, 2026 17:39

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-foundation branch from 60522c1 to bb9c235Compare August 30, 2026 00:49

Navigation Menu

Uh oh!

fix: TensorRT-RTX Turing (SM 7.5) capability guards and compute-capability targeting - #4643

fix: TensorRT-RTX Turing (SM 7.5) capability guards and compute-capability targeting#4643
tp5uiuc wants to merge 3 commits into
mainfrom
tp5uiuc/trtrtx-turing-foundation

tp5uiuc commented Aug 29, 2026

Uh oh!

tp5uiuc commented Aug 29, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

tp5uiucAug 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		# Bind to a typed local: the imported helper is untyped from mypy's view here,
		# and returning it directly trips --strict's no-any-return.

		# An engine built for one set of compute capabilities cannot be reused for a
		# compile targeting a different set -- doing so would silently reintroduce ops the

Uh oh!

Conversation

tp5uiuc commented Aug 29, 2026

Uh oh!

tp5uiuc commented Aug 29, 2026

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant