Skip to content

fix: guard the GEMM the cdist converter emits on Turing - #4645

Open
tp5uiuc wants to merge 1 commit into
tp5uiuc/trtrtx-turing-conv3d-padfoldfrom
tp5uiuc/trtrtx-turing-cdist-gemm
Open

fix: guard the GEMM the cdist converter emits on Turing#4645
tp5uiuc wants to merge 1 commit into
tp5uiuc/trtrtx-turing-conv3d-padfoldfrom
tp5uiuc/trtrtx-turing-cdist-gemm

Conversation

@tp5uiuc

@tp5uiuc tp5uiuc commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

What — Adds cdist_forward_capability_validator, reproducing the converter's own branch
condition, so an aten._cdist_forward that would emit a matrix-multiply layer falls back to PyTorch
on Turing.

WhyThis closes zero failing tests, and that is the point. impl.normalization.cdist_forward
computes p == 2 with a matrix-multiply layer when compute_mode == 1, or when it is 0/absent and
either operand has more than 25 rows. The GEMM is emitted inside the converter, so the graph holds
a single _cdist_forward node and no mm/bmm for gemm_capability_validator to reject.
The covering tests were skipped rather than the converter guarded, so nothing was ever red while the
conversion path stayed open. A skip protects CI and does not protect a caller.

How — A validator reproducing the converter's own p / compute_mode / row-count branch,
so the guard fires exactly when the matmul layer would be emitted. Deliberately not keyed on
dtype, unlike gemm_capability_validator: measured on a T4, the fused Matmul_MUL_SUB_SQRT_ pattern
fails for FP16 operands too, under enabled_precisions {f32}, {f16} and {f16,f32} alike, even
though a bare FP16 matmul and an FP16 nn.Linear both run there. What predicts the failure is whether
the matmul layer is emitted, not its precision.

Testing — Confirmation sweep, both arms; test_cdist_aten.py and the new guard cases. No status
change on the L40S.

Cost / Gotchas — The 25-row threshold is load-bearing, not an optimisation: at or below it no GEMM
is emitted and cdist runs correctly on Turing in both dtypes. Rejecting those anyway would regress,
because PyTorch's cdist_cuda kernel has no Half implementation, so the fallback would raise where
TensorRT-RTX succeeds today.

🤖 Generated with Claude Code

@github-actions github-actions Bot added component: tests Issues re: Tests component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 29, 2026
@meta-cla meta-cla Bot added the cla signed label Aug 29, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 29, 2026 17:54
@tp5uiuc
tp5uiuc changed the base branch from tp5uiuc/trtrtx-turing-foundation to tp5uiuc/trtrtx-turing-conv3d-padfold August 29, 2026 18:07
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from 750021c to 4d25c1d Compare August 29, 2026 18:09

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/models/test_turing_capability_guards.py	2026-08-29 18:09:20.356851+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/models/test_turing_capability_guards.py	2026-08-29 18:09:54.488199+00:00
@@ -113,12 +113,10 @@
def _cdist_inputs(rows1, rows2, dtype=torch.float32):
    return (
        torch.randn(4, rows1, 5, dtype=dtype).cuda(),
        torch.randn(4, rows2, 5, dtype=dtype).cuda(),
    )
-
-


@unittest.skipIf(
    not ENABLED_FEATURES.tensorrt_rtx,
    "Turing capability guards only apply to TensorRT-RTX",

@tp5uiuc tp5uiuc self-assigned this Aug 29, 2026
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from 4d25c1d to 3628e48 Compare August 29, 2026 18:17
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from 3628e48 to c3412ff Compare August 29, 2026 19:05
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from c3412ff to eb3ef0d Compare August 29, 2026 23:01
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from eb3ef0d to aa9e28f Compare August 30, 2026 00:49
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from aa9e28f to 556ef8d Compare August 30, 2026 00:52
Comment thread py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py Outdated
Comment thread py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py Outdated
Comment thread py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py Outdated
Comment thread py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py Outdated
Comment thread py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py Outdated
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from 556ef8d to ea3d105 Compare August 30, 2026 00:58
@github-actions
github-actions Bot requested a review from narendasan August 30, 2026 04:15
impl.normalization.cdist_forward computes p == 2 with a matrix-multiply layer for
compute_mode 1, or 0/absent with an operand above the row threshold. That GEMM is
emitted inside the converter, so the graph holds a single _cdist_forward node and
no mm/bmm for gemm_capability_validator to reject, and on Turing TensorRT-RTX
fails cuDNN graph compilation with "No valid engine configs for
Matmul_MUL_SUB_SQRT_", leaving a null execution context.

Zero tests were failing: PR #4546 skipped the covering tests rather than guarding
the converter, so nothing was red while the conversion path stayed open. A skip
protects CI; it does not protect a caller.

Which arguments emit a GEMM is decided by cdist_emits_matmul, the predicate
cdist_forward itself uses, so the guard cannot drift from the branch it mirrors.
Unlike gemm_capability_validator this is deliberately not dtype-keyed: measured on
a T4 the fused pattern fails for FP16 operands too, so what predicts failure is
whether the matmul layer is emitted, not its precision. The row threshold is
load-bearing rather than an optimisation -- below it no GEMM is emitted, and
rejecting anyway would regress, since PyTorch's cdist_cuda has no Half kernel.

One behaviour note: where exactly one operand's row count is statically known and
exceeds the threshold, the shared predicate rejects, whereas reading the shapes
directly would have failed open. Rejecting is correct -- the converter's condition
is an or, so a single known operand above the threshold already decides it. The
case is unreachable today because the converter is not registered
supports_dynamic_shapes, so the partitioner refuses a partially dynamic cdist
first.
@tp5uiuc
tp5uiuc force-pushed the tp5uiuc/trtrtx-turing-cdist-gemm branch from ea3d105 to a0c3245 Compare August 30, 2026 18:37
@tp5uiuc

tp5uiuc commented Aug 30, 2026

Copy link
Copy Markdown
Collaborator Author

CI summary

failing check cause
standard / test (dynamo-models, standard, l2) test_cross_runtime_serde::test_save_cpp_load_python and ::test_save_python_load_cpp — Windows only
standard / test (dynamo-models-critical, standard, l1) test_hf_gqa_model::test_dynamic_head_dim_with_hf_model[False-dtype0]Tensor-likes are not close
executorch-runtime-test OSError: libcurand.so.10
gate ×2 aggregators

Suites ran to completion: 2 failed, 150 passed, 61 skipped, 4 xpassed and 1 failed, 52 passed.

Why these are not from this PR: the same test_cross_runtime_serde pair fails identically on #4643, #4644, #4645, #4646, #4647 and #4648, whose change sets are disjoint — e.g. #4648 changes _exporter.py (the save path) and #4647 changes neither that nor anything else in common, yet both fail the same two tests with byte-identical error text. executorch-runtime-test reproduces on main nightlies and on unrelated PRs (#4638, #4642). The gate entries are aggregators of the above. test_hf_gqa_model also fails on #4643 and #4647 and has no overlap with anything this PR changes.

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

Labels

ci: full cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant