fix: guard the FP32 GEMMs that reach TensorRT-RTX through linear and attention by tp5uiuc · Pull Request #4646 · pytorch/TensorRT

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

What — Registers gemm_capability_validator on aten.linear.default and on the SDPA converters,
and follows that guard through the four test sites it changes.

Why — The GEMM guard covered matmul/mm/bmm/dot/mv/addmm, but not aten.linear.default
— and Torch-TensorRT deliberately disables the linear -> addmm decomposition, so the guarded node
never appears. Attention is converted as a unit, so no mm/bmm node is ever created either. Both
reach TensorRT-RTX on Turing and fail:

cuDNN graph compilation failed: No valid engine configs for Matmul_ -> null execution context

and on dynamic shapes the Myelin error is logged at enqueueV3 but not raised — 3 of the 19
affected cases returned silently wrong values.

How — Register the validator on aten.linear.default and on the three attention converters,
then follow it everywhere it lands. The resource-partitioning test's hard-coded accelerated and
non-accelerated counts no longer hold once the model's trailing FP32 nn.Linear is pushed into a
PyTorch block. Three model tests reach an entry point that emits one engine for the whole program and
has no partitioner to fall back to, so the rejection raises instead of producing a PyTorch block —
those run their model in FP16 on Turing, keyed off the same predicate the validator uses so the two
cannot drift.

Testing — Confirmation sweep, both arms. 18 failures closed directly; the guard additionally costs
10 tests (2 in partitioning/, 8 in models/) which this PR carries the fixes for. Only 4 of
those 10 showed as pass -> fail — the other 6 were already failing under a different cause, so a
status diff alone does not find them. Neutralising the guard at runtime brings back the null execution
context and the silent all-zeros output, confirming it is correct and necessary. No status change on
the L40S.

Cost / Gotchas

Coverage: wherever the guard fires, the Turing arm exercises less than the L40S arm — measured
at ~10% of the streamable weights for the weight-streaming model, which is FP32 nn.Linear-heavy.
The weight-stripping and refit tests now run in FP16 on Turing and FP32 elsewhere, so the two
arms are no longer testing the same dtype. The assertions (engine sizes, stripping, refit) are
equally meaningful in FP16 and the network is strongly typed, but it is a real asymmetry.

Followups — 5 models/ cosine-similarity failures on Turing (0.9862–0.9872 against a 0.99
threshold) are not addressed here. They compile fully to TensorRT with no capability fallback
logged, which rules the guards out, and they were already failing before this series. Not diagnosed.

🤖 Generated with Claude Code

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

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.

…attention TensorRT-RTX cannot serve FP32 GEMMs on Turing (SM 7.5). The capability validator that enforces this is attached to the explicit GEMM aten targets -- matmul, mm, bmm, mv, dot and addmm -- which misses the two paths that carry a GEMM without ever producing one of those nodes: * aten.linear.default is registered with no capability validator, and its decomposition into addmm is deliberately disabled, so the guarded addmm node never appears in the graph to be rejected. * Attention is converted as a single fused subgraph. Its matmuls are emitted inside the converter, so there is no mm/bmm node for the GEMM validator to inspect. On Turing both paths reach TensorRT-RTX and fail cuDNN graph compilation with "No valid engine configs for Matmul_ADD_" / "Matmul_MUL_". Under static shapes that surfaces as a null execution context. Under dynamic shapes it does not surface at all: the error is logged from IExecutionContext::enqueueV3, execution returns normally, and the caller receives an all-zero tensor of the correct shape and dtype. Measured on a T4, an FP32 nn.Linear and an FP32 fused attention block with a dynamic batch dimension both returned 100% zeros, cosine similarity 0.0 against eager, and raised nothing. That silent case covers essentially every FP32 transformer on Turing, and closing it is the point of this change. * Register aten.linear.default with the existing gemm_capability_validator. * Add attention_capability_validator and call it from the scaled-dot-product, flash-attention and efficient-attention validators; the cuDNN-attention validator inherits it through the efficient one it already delegates to. The rejected dtypes live in a single module-level tuple. * Key both guards on FP32 only. FP16 GEMMs and FP16 attention run correctly on Turing and must keep running on TensorRT. * Skip the FP32 converter unit tests for linear and attention on Turing. DispatchTestCase has no PyTorch-fallback path, so once a validator rejects an op those tests raise UnsupportedOperatorException instead of falling back; without the skips the guard would simply trade one failure for another. * Extend the Turing guard tests with linear and fused-attention coverage, including the dynamic-shape cases that previously returned zeros. Most of it runs under target_compute_capabilities=[(7, 5)], so it exercises the guards on any GPU. Testing (T4 / SM 7.5 and L40S / SM 8.9, driver 595.58.03, identical stacks): * The 18 previously-failing tests all pass or skip on Turing. The two that run through real export flows -- the force-causal efficient-attention lowering test and the default weight-streaming runtime test -- now pass rather than skip, confirming the fallback produces correct numbers. * The three dynamic-shape cases that previously returned zeros (aten.linear, and fused attention with and without a float mask) go from 100% zero output and cosine 0.0 to zero maximum absolute error against eager, with no TensorRT engine built for the guarded subgraph. * Full lowering/, runtime/ and conversion/ sweeps on both arms: no test regressed on either. On Turing, 6 tests go fail->pass and 16 go fail->skip; 4 more go pass->skip, all FP32 shapes that happened to work there and that a dtype-keyed validator cannot distinguish. On the non-Turing control every one of the 2603 tests keeps its previous status, confirming the guards are inert off Turing.

…M guard ce480a0 registered torch.ops.aten.linear.default with gemm_capability_validator. Every model in tests/py/dynamo/partitioning/test_001_resource_partitioning.py ends in an FP32 nn.Linear(1024*56*56, 10), so on Turing (SM 7.5) that guard rejects it, capability partitioning routes it to a PyTorch block, and the hard-coded acc/gpu splits no longer hold. Three tests were green on both arms at cf10a41 and red on the T4 after, with no numerical error, no crash and no TensorRT message anywhere in the run -- every failure is a partition-count expectation. Attribution was proven by three-way runtime neutralisation of aten.linear.default, with aten.convolution.default as the negative control. Only two of the three are guard failures. The third is a cascade, and the control that separates them is running the global test alone: on the T4, against the unmodified tree, test_resource_partitioning_with_global_capability_partitioning PASSES on its own with its original == 4. It only fails when the atomic-subgraphs test runs first and fails, because that test's ATOMIC_SUBGRAPHS.remove((ReLUConv, (), True)) was the statement after the assertion that fails, so it never ran and ReLUConv leaked into the next test. A leaked atomic subgraph stops the resource partitioner splitting _run_on_acc_0, which is exactly what that test's own docstring describes, and the count drops from 4 to 3. The conftest fixture restores the converter registry but not ATOMIC_SUBGRAPHS. * Derive the expected non-accelerated count in the two capability-partitioning tests from trt_rtx_targets_turing() -- the same predicate the validator keys on, so the expectation cannot drift from the guard -- rather than hard-coding 2. Reusing the library helper also avoids adding a fourth copy of the SM 7.5 check. * Unregister ReLUConv through addCleanup instead of a trailing statement, so a failed assertion can no longer leak it into the next test. * Leave test_resource_partitioning_with_global_capability_partitioning alone. Once the leak is fixed it passes on Turing unchanged. Its 4 blocks are composed differently there -- the Linear is a PyTorch block and _run_on_acc_0 splits in two, rather than the Linear being its own accelerated block -- so the count matches for a different reason. Asserting that composition is beyond restoring the pre-branch signal and is not done here. Making the counts capability-aware rather than skipping is deliberate. Cause G skipped because DispatchTestCase hands the graph straight to TRTInterpreter, so a rejected converter raises and the test cannot run at all on Turing. The partitioner path has a PyTorch fallback, so these tests do run and do produce a correct partitioning; only the number was stale. Skipping would switch off the only tests covering resource partitioning composed with a capability fallback, on the one architecture where that fallback happens, and would take partitioning/'s contribution to this branch's Turing skips from 0 to 3. The cost is that the Turing arm now asserts a different number from the non-Turing arm, so the fallback branch is only exercised on Turing hardware, not in CI. Testing (T4 / SM 7.5 ipp1-2023 and L40S / SM 8.9 a1u1g-mil-0589, driver 595.58.03, identical stacks; before and after measured on the same node and container): * partitioning/ on the T4: 21 passed / 3 failed / 0 skipped -> 24 passed / 0 failed / 0 skipped, 24 collected. On the L40S: 24 / 0 / 0 both before and after. Reconciled by t3877f-check.py against four independent sources on each arm: RESULT: COMPLETE. * Exactly the three target tests change status, all fail -> pass, and only on the T4. Every other test in partitioning/ keeps its status on both arms; the L40S per-test status list is byte-identical before and after. Cross-arm diff after the change: 0 Turing-specific failures. * The guard is still firing. Re-running the same neutralisation control on the T4 after the change inverts it, as it must: none -> 3 passed (guard on, tests expect the Turing split), linear -> 2 failed (guard off, split reverts, the predicate still reports Turing), conv -> 3 passed. The linear run also shows the leak is fixed: the atomic-subgraphs test fails there and the global test still passes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

… guard ce480a0 registered torch.ops.aten.linear.default with gemm_capability_validator, which rejects an FP32 GEMM when Turing (SM 7.5) is a build target. Four tests in tests/py/dynamo/models/ reach that rejection on a code path that converts a whole graph module directly, with no partitioner to route the node into a PyTorch block, so on a T4 all four raise UnsupportedOperatorException: Conversion of function torch._ops.aten.aten::linear not currently supported! Two are visible pass->fail regressions. The other two were already failing at cf10a41 for cause D, which is now fixed, so the guard replaced their reason under an unchanged status and no status diff flags them. Attribution is the three-way runtime neutralisation of aten.linear.default over the 16 Turing-specific models/ failures, with aten.convolution.default as the negative control: none 16 failed, linear 10 failed and 6 passed (these 4 among them), conv 16 failed 0 passed. The guard is correct and is not touched here. Neutralise it and these paths build an engine that on SM 7.5 either fails to create an execution context or returns a silent all-zero tensor of the right shape and dtype -- the two failure modes ce480a0 exists to prevent. Two of the entry points have no fallback by design: convert_exported_program_to_serialized_trt_engine emits one engine for the whole program, and construct_refit_mapping interprets a module to build a weight map. * The three convert_exported_program_to_serialized_trt_engine tests run their model in FP16 when trt_rtx_targets_turing() -- the same predicate the validator keys on, so the expectation cannot drift from the guard, and the same library helper cause K reused rather than a fourth copy of the SM 7.5 check. None of the three is about FP32: two are about kwarg_inputs plumbing and one compares a weight-stripped against a weight-included engine size, a comparison that is just as real in FP16. The guard keys on operand dtype only and Turing has FP16 GEMM hardware; no enabled_precisions change is needed because the network is strongly typed. * test_model_refit.py::test_mapping partitions before mapping. construct_refit_mapping interprets whatever module it is handed and has no partitioner, so it has to be handed the subgraph the engine was built from. refit_module_weights always partitions first and maps each accelerated submodule; this test was the only caller in the tree passing a whole un-partitioned module, and that only worked because resnet18 happens to be fully convertible on SM 8.9. On Turing its fc GEMM is rejected, so the engine it compares against came from _run_on_acc_0 without the Linear while the mapping was built from a graph with it -- it was mapping the wrong graph, and would have been even had it not raised. This half is not architecture-conditional and changes nothing on SM 8.9, where _run_on_acc_0 is the whole graph. Making the tests follow the capability rather than skipping is deliberate, and is what cause K did for the partition counts. Cause G skipped because DispatchTestCase cannot run at all on Turing once a converter is rejected; here the tests can run, they were just asking for a dtype the target cannot serve. Skipping instead would take models/'s contribution to this branch's Turing skips from 0 to 4, two of them currently green, on a branch that already switches off 105 tests on Turing. The cost is that the Turing arm now exercises those three paths in FP16 while the non-Turing arm keeps exercising them in FP32, so the FP32 no-fallback path is only covered off Turing -- it cannot be covered on Turing, where it correctly raises. test_mapping loses nothing: it stays FP32 on both arms and now maps the subgraph the engine contains. Testing (T4 / SM 7.5 ipp1-2023 and L40S / SM 8.9 a1u1g-mil-0572, driver 595.58.03, identical stacks, -n 1, full models/ module runs on both arms): * models/ on the T4: 225 passed / 19 failed / 25 skipped -> 229 passed / 15 failed / 25 skipped, 269 collected (1:00:22). On the L40S: 238 / 4 / 27 both before and after (17:53). Reconciled by t3877f-check.py against four independent sources on each arm -- the summary line, the collected count, the junit XML and the streamed progress lines: RESULT: COMPLETE on both. * Exactly the four target tests change status, all fail -> pass, and only on the T4. A line-by-line diff of the per-test status list against the pre-change baseline shows those four lines and nothing else on the T4, and is byte-identical on the L40S. * The other 15 T4 failures are unchanged and none is mine: 12 remain Turing-specific (causes M, N, O, P, Q) and 3 fail on both arms (view_as_real). Cross-arm comparison after the change lists 0 of the 4 as Turing-specific. * The guard is still firing. Re-running the exact pre-change FP32 model through convert_exported_program_to_serialized_trt_engine on the T4 still raises UnsupportedOperatorException on aten.linear.default; the FP16 form now builds a 2,619,588-byte engine; and the same FP32 model through torchtrt.dynamo.compile still partitions to ['_run_on_acc_0', '_run_on_gpu_1'], i.e. the guard fires there too and only differs in having somewhere to send the node. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tp5uiuc · 2026-08-30T00:55:25Z

+ for that guard to see. On Turing the result is a cuDNN graph-compilation failure --
+ a null execution context under static shapes, a silently wrong result under dynamic
+ ones. Only the q/k/v dtypes matter, so FP16 attention keeps running on TensorRT.
+ """


On Turing TRT-RTX this results in failure. Thats it no need to say more. Only the q/k/v dtypes matter, so FP16 attention keeps running on TensorRT. this can be kept.

tp5uiuc · 2026-08-30T00:55:50Z

+# Attention is converted as a single fused subgraph, so the matmuls it performs never
+# appear as nodes in the graph and ``gemm_capability_validator`` never sees them; the
+# dtypes that GEMM cannot serve have to be listed again here.
+_TURING_UNSUPPORTED_ATTENTION_DTYPES = (torch.float32,)


Push it inside the function

tp5uiuc · 2026-08-30T00:56:19Z

+ val = arg.meta.get("val") if hasattr(arg, "meta") else None
+ if val is not None and (
+ getattr(val, "dtype", None) in _TURING_UNSUPPORTED_ATTENTION_DTYPES
+ ):


Same readability concernts as #4643 etc.

ce480a0 registered torch.ops.aten.linear.default with gemm_capability_validator, which rejects an FP32 GEMM when Turing (SM 7.5) is a build target. Four tests in tests/py/dynamo/models/test_weight_stripped_engine.py compile resnet18 in FP32 with strip_engine_weights=True and assert the module's output is all zeros. On a T4 the guard rejects resnet18's trailing fc, the partitioner routes it into a PyTorch block, and a PyTorch block holds the real nn.Linear parameters. strip_engine_weights=True strips the *engine's* weights and cannot touch those, so the stripped engine emits zeros and the PyTorch fc adds its real bias on top: the module's output is that bias -- non-zero, and identical in every row of the batch. test_compile_weight_stripped_engine tensor(123.7316) != 0 test_weight_stripped_engine_results tensor(-0.0001) != 0 test_two_TRTRuntime_in_refitting tensor(-0.0001) != 0 test_refit_weight_stripped_engine_multiple_times tensor(-0.0076) != 0 The guard is correct and is not touched here. Neutralising aten.linear.default makes the all-zeros assertion pass and then each of the four fails on a *later* assertion instead -- a null execution context, a cosine-similarity mismatch, or a refitted engine that is still all zeros. Those are exactly the two failure modes ce480a0 exists to prevent, so a status-only control would have read as "the guard is not responsible" while the reason changed completely. The fix is the same one cause L (6470ebb) used for test_weight_stripped_engine_sizes in this same file: run the model in FP16 when trt_rtx_targets_turing() -- the same predicate the validator keys on, so the test expectation cannot drift from the guard. The guard keys on operand dtype only and Turing has FP16 GEMM hardware, so in FP16 the whole graph converts, there is no PyTorch block, and the stripped engine's output is the exact zeros the tests assert. No enabled_precisions change is needed because the network is strongly typed. Cause L spelled that dtype choice inline. Rather than add four more copies, this hoists it to a module-level _turing_safe_gemm_dtype() helper whose docstring records both ways the fc rejection surfaces in this file -- the no-fallback path, where convert_exported_program_to_serialized_trt_engine raises UnsupportedOperatorException, and the partitioned path above -- and switches cause L's test to the helper too, so the file carries one copy of the rule instead of five. That half is behaviourally identical on both arms. Nothing outside this file changes; in particular aten_ops_converters.py, which cause C/E owns, is untouched. What each test still asserts on Turing, unchanged from before: * test_compile_weight_stripped_engine -- a torch_trt.compile(ir="dynamo") build with strip_engine_weights=True yields a module whose output is all zeros. * test_weight_stripped_engine_results -- on a dynamic batch dim: stripped output is all zeros, refitting with the same weights makes it non-zero, and the refitted output matches a separately torch.compile'd weight-included engine above COSINE_THRESHOLD. * test_two_TRTRuntime_in_refitting -- over two independent compile+refit cycles in one process: stripped output all zeros, and the refitted output matches eager PyTorch above COSINE_THRESHOLD. Both iterations, both assertions. * test_refit_weight_stripped_engine_multiple_times -- stripped output all zeros, a first refit yields non-zero, a second refit onto the already-refitted engine (the INCLUDE_REFIT path) yields non-zero at a different shape, and that matches a torch.compile'd weight-included engine above COSINE_THRESHOLD. None of the four loses an assertion, a code path or a compile entry point. The cost is that the Turing arm exercises weight stripping, refit and multiple runtimes in FP16 while the non-Turing arm keeps exercising them in FP32, so the FP32 form of these paths is covered only off Turing -- it cannot be covered on Turing, where the guard correctly refuses the GEMM. Skipping instead would take models/'s contribution to this branch's Turing skips from 0 to 4 on a branch that already switches off 105 tests on Turing, and cause L deliberately declined to do that in this very file. Nothing about weight stripping is dtype-specific. Testing (T4 / SM 7.5 ipp1-2023 and L40S / SM 8.9 a1u1g-mil-0572, driver 595.58.03, identical stacks, -n 1, full models/ runs on both arms with --ignore=models/test_hf_gqa_model.py, 266 collected): * T4: 224 passed / 13 failed / 25 skipped / 4 xpassed -> 228 passed / 9 failed / 25 skipped / 4 xpassed (48:36). L40S: 231 / 4 / 27 / 4 both before and after (14:33). Reconciled by t3877f-check.py against the summary line, the collected count, the junit XML and the streamed progress lines: RESULT: COMPLETE on both. * Exactly the four target tests change status, all fail -> pass, and only on the T4. A per-test status diff against the pre-change baseline shows those four lines and nothing else on the T4, and reports CHANGED: 0 on the L40S. * The 9 remaining T4 failures are unchanged and none is mine: 6 remain Turing-specific (5 cosine-similarity, cause O; 1 bert cpu_offload, cause P) and 3 fail on both arms (view_as_real). The other 8 tests in this file keep their status on both arms, including the 2 that skip and cause L's test_weight_stripped_engine_sizes. * The guard is still firing. After the change, the unmodified FP32 resnet18 through torch_trt.dynamo.compile on the T4 still partitions to ['_run_on_acc_0', '_run_on_gpu_1'] and still returns a non-zero stripped output (-0.000119), and the same FP32 model through convert_exported_program_to_serialized_trt_engine still raises UnsupportedOperatorException. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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

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

tp5uiuc mentioned this pull request Aug 29, 2026
fix: guard the GEMM the cdist converter emits on Turing #4645
Open

tp5uiuc changed the base branch from tp5uiuc/trtrtx-turing-foundation to tp5uiuc/trtrtx-turing-cdist-gemmAugust 29, 2026 18:07

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from 598a7f3 to c5c8860Compare August 29, 2026 18:09

tp5uiuc self-assigned this Aug 29, 2026

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from c5c8860 to a08d91bCompare August 29, 2026 18:17

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from a08d91b to b75a683Compare August 29, 2026 19:05

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from b75a683 to 3d3e181Compare August 29, 2026 23:01

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from 3d3e181 to 0fd0559Compare August 30, 2026 00:49

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from 0fd0559 to 5d31a6bCompare August 30, 2026 00:53

tp5uiucand others added 3 commits August 29, 2026 17:56

tp5uiuc commented Aug 30, 2026
View reviewed changes

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from 5d31a6b to fe5ee5bCompare August 30, 2026 00:58

tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-fp32-gemm-linear branch from fe5ee5b to c85d3dbCompare August 30, 2026 01:00

Navigation Menu

Uh oh!

fix: guard the FP32 GEMMs that reach TensorRT-RTX through linear and attention - #4646

fix: guard the FP32 GEMMs that reach TensorRT-RTX through linear and attention#4646
tp5uiuc wants to merge 4 commits into
tp5uiuc/trtrtx-turing-cdist-gemmfrom
tp5uiuc/trtrtx-turing-fp32-gemm-linear

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

tp5uiuc commented Aug 29, 2026

Uh oh!

tp5uiuc commented Aug 29, 2026

Uh oh!

tp5uiucAug 30, 2026

Choose a reason for hiding this comment

Uh oh!

tp5uiucAug 30, 2026

Choose a reason for hiding this comment

Uh oh!

tp5uiucAug 30, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant