fix: guard the GEMM the cdist converter emits on Turing - #4645
Conversation
750021c to
4d25c1d
Compare
There was a problem hiding this comment.
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",4d25c1d to
3628e48
Compare
3628e48 to
c3412ff
Compare
c3412ff to
eb3ef0d
Compare
eb3ef0d to
aa9e28f
Compare
aa9e28f to
556ef8d
Compare
556ef8d to
ea3d105
Compare
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.
ea3d105 to
a0c3245
Compare
CI summary
Suites ran to completion: Why these are not from this PR: the same |
What — Adds
cdist_forward_capability_validator, reproducing the converter's own branchcondition, so an
aten._cdist_forwardthat would emit a matrix-multiply layer falls back to PyTorchon Turing.
Why — This closes zero failing tests, and that is the point.
impl.normalization.cdist_forwardcomputes
p == 2with a matrix-multiply layer whencompute_mode == 1, or when it is 0/absent andeither operand has more than 25 rows. The GEMM is emitted inside the converter, so the graph holds
a single
_cdist_forwardnode and nomm/bmmforgemm_capability_validatorto 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 fusedMatmul_MUL_SUB_SQRT_patternfails for FP16 operands too, under
enabled_precisions{f32},{f16}and{f16,f32}alike, eventhough a bare FP16 matmul and an FP16
nn.Linearboth run there. What predicts the failure is whetherthe matmul layer is emitted, not its precision.
Testing — Confirmation sweep, both arms;
test_cdist_aten.pyand the new guard cases. No statuschange 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_cudakernel has no Half implementation, so the fallback would raise whereTensorRT-RTX succeeds today.
🤖 Generated with Claude Code