fix: guard the GEMM the cdist converter emits on Turing

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -3369,7 +3369,55 @@ def aten_ops_convolution(
         )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten._cdist_forward.default)
+# Above this many rows in either operand, the p == 2 path of
+# impl.normalization.cdist_forward switches from a broadcast-subtract to a matrix
+# multiply. Kept deliberately in sync with the threshold in that converter.
+# aten._cdist_forward(x1, x2, p, compute_mode)
+_CDIST_ARG_P = 2
+_CDIST_ARG_COMPUTE_MODE = 3
+
+
+def cdist_forward_capability_validator(
+    node: Node, settings: Optional[CompilationSettings] = None
+) -> bool:
+    """Reject the cdist variants whose converter emits a GEMM, on Turing (SM 7.5).
+
+    The GEMM is emitted *inside* the converter, so the graph holds one
+    ``_cdist_forward`` node and no ``mm``/``bmm`` for ``gemm_capability_validator`` to
+    reject, and TensorRT-RTX then fails. Which arguments emit one is decided by
+    ``cdist_emits_matmul``, the converter's own predicate, so the two cannot drift.
+    """
+    if not trt_rtx_targets_turing(settings):
+        return True
+
+    def operand_rows(operand: Argument) -> Optional[int]:
+        val = operand.meta.get("val") if hasattr(operand, "meta") else None
+        shape = getattr(val, "shape", None)
+        if shape is None or len(shape) < 2:
+            return None
+        rows = shape[-2]
+        return rows if isinstance(rows, int) else None
+
+    operands = node.args[:2]  # x1, x2
+    if not impl.normalization.ops.cdist_emits_matmul(
+        args_bounds_check(node.args, _CDIST_ARG_P, replacement=2.0),
+        args_bounds_check(node.args, _CDIST_ARG_COMPUTE_MODE, replacement=None),
+        [operand_rows(operand) for operand in operands],
+    ):
+        return True
+
+    _LOGGER.debug(
+        "cdist '%s' computes p=2 as a matrix multiply, which is not supported on "
+        "TensorRT-RTX for Turing (SM 7.5). Falling back to PyTorch.",
+        node.name,
+    )
+    return False
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten._cdist_forward.default,
+    capability_validator=cdist_forward_capability_validator,
+)
 def aten_ops_cdist_forward(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/models/test_turing_capability_guards.py

@@ -97,6 +97,30 @@ def forward(self, x):
         return self.conv(F.pad(x, (0, 1, 0, 1)))
 
 
+class Cdist(nn.Module):
+    """A GEMM the explicit GEMM guards never see, emitted inside the cdist converter.
+
+    The op has to be called directly: torch.cdist does not reach this converter on the
+    GEMM path, because ATen's cdist calls _euclidean_dist there, which lowers to
+    aten.matmul and is guarded already. Only _cdist_forward's own p == 2 branch emits a
+    matmul layer that no mm/bmm node stands for.
+    """
+
+    def __init__(self, p=2.0, compute_mode=None):
+        super().__init__()
+        self.p, self.compute_mode = p, compute_mode
+
+    def forward(self, x1, x2):
+        return torch.ops.aten._cdist_forward.default(x1, x2, self.p, self.compute_mode)
+
+
+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",
@@ -205,6 +229,59 @@ def test_declared_turing_target_keeps_transposed_conv3d_on_trt(self):
         compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
         self.assertGreater(_trt_submodule_count(compiled), 0)
 
+    def test_declared_turing_target_falls_back_cdist_p2_compute_mode_1(self):
+        # compute_mode=1 is "always use the matrix multiply", whatever the row count.
+        mod = Cdist(compute_mode=1)
+        inputs = _cdist_inputs(6, 8)
+        compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
+        self.assertEqual(_trt_submodule_count(compiled), 0)
+        self._assert_matches_eager(mod, compiled, inputs)
+
+    def test_declared_turing_target_falls_back_cdist_p2_default_compute_mode(self):
+        # An absent compute_mode means "use the matrix multiply if either operand has
+        # more than 25 rows". The validator has to normalise None to 0 as the converter does.
+        mod = Cdist()
+        inputs = _cdist_inputs(35, 45)
+        compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
+        self.assertEqual(_trt_submodule_count(compiled), 0)
+        self._assert_matches_eager(mod, compiled, inputs)
+
+    def test_declared_turing_target_falls_back_fp16_cdist_p2(self):
+        # Unlike every other guard here, this one is not FP32-only: on a T4 the fused
+        # Matmul_MUL_SUB_SQRT_ pattern fails for FP16 operands too, under every
+        # enabled_precisions setting, even though a bare FP16 GEMM runs there fine.
+        mod = Cdist(compute_mode=1)
+        inputs = _cdist_inputs(6, 8, dtype=torch.half)
+        compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
+        self.assertEqual(_trt_submodule_count(compiled), 0)
+        self._assert_matches_eager(mod, compiled, inputs)
+
+    def test_declared_turing_target_keeps_cdist_p2_compute_mode_2_on_trt(self):
+        # compute_mode=2 is "never use the matrix multiply", so there is no GEMM to reject.
+        mod = Cdist(compute_mode=2)
+        inputs = _cdist_inputs(35, 45)
+        compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
+        self.assertGreater(_trt_submodule_count(compiled), 0)
+        self._assert_matches_eager(mod, compiled, inputs)
+
+    def test_declared_turing_target_keeps_cdist_p2_small_rows_on_trt(self):
+        # At or below 25 rows the default compute_mode emits no GEMM. Rejecting these
+        # anyway would be worse than useless: PyTorch's cdist_cuda kernel has no Half
+        # implementation, so for FP16 the fallback raises where TensorRT-RTX succeeds.
+        mod = Cdist(compute_mode=0)
+        inputs = _cdist_inputs(6, 8)
+        compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
+        self.assertGreater(_trt_submodule_count(compiled), 0)
+        self._assert_matches_eager(mod, compiled, inputs)
+
+    def test_declared_turing_target_keeps_cdist_p1_on_trt(self):
+        # Only p == 2 has a matrix-multiply path at all.
+        mod = Cdist(p=1.0, compute_mode=0)
+        inputs = _cdist_inputs(35, 45)
+        compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING])
+        self.assertGreater(_trt_submodule_count(compiled), 0)
+        self._assert_matches_eager(mod, compiled, inputs)
+
     @unittest.skipIf(
         _is_turing(), "on Turing the native path is already guarded; see the SM75 tests"
     )
@@ -387,6 +464,19 @@ def test_conv3d_falls_back(self):
         compiled = self._compile(mod, inputs)
         self.assertEqual(_trt_submodule_count(compiled), 0)
 
+    def test_cdist_p2_falls_back(self):
+        # Unguarded this built an engine whose createExecutionContext() then returned
+        # nullptr: "cuDNN graph compilation failed: No valid engine configs for
+        # Matmul_MUL_SUB_SQRT_".
+        mod = Cdist(compute_mode=1)
+        inputs = _cdist_inputs(6, 8)
+        compiled = self._compile(mod, inputs)
+        self.assertEqual(_trt_submodule_count(compiled), 0)
+        with torch.no_grad():
+            ref = mod.eval().cuda()(*inputs)
+            out = compiled(*inputs)
+        self.assertEqual((ref - out).abs().max().item(), 0.0)
+
     def test_bfloat16_falls_back_without_crashing(self):
         # Unguarded, compiling bfloat16 for SM 7.5 segfaulted the process.
         class Add(nn.Module):