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

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -4519,7 +4519,11 @@ def aten_ops_nonzero(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.linear.default, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.linear.default,
+    capability_validator=gemm_capability_validator,
+    supports_dynamic_shapes=True,
+)
 def aten_ops_linear(
     ctx: ConversionContext,
     target: Target,
@@ -4538,9 +4542,43 @@ def aten_ops_linear(
     )
 
 
+def attention_capability_validator(
+    node: Node, settings: Optional[CompilationSettings] = None
+) -> bool:
+    """Reject fused attention TensorRT-RTX cannot serve when Turing (SM 7.5) is a target.
+
+    A fused attention op carries the same FP32 GEMMs ``gemm_capability_validator``
+    rejects, but inside the converter, so the graph holds one node and no ``mm``/``bmm``
+    for that guard to see. On Turing TRT-RTX this results in failure. Only the q/k/v
+    dtypes matter, so FP16 attention keeps running on TensorRT.
+    """
+    if not trt_rtx_targets_turing(settings):
+        return True
+
+    unsupported_dtypes = (torch.float32,)
+
+    def is_unsupported(operand: Argument) -> bool:
+        val = operand.meta.get("val") if hasattr(operand, "meta") else None
+        return bool(getattr(val, "dtype", None) in unsupported_dtypes)
+
+    qkv = node.args[:3]  # query, key, value
+    if any(map(is_unsupported, qkv)):
+        _LOGGER.debug(
+            "Attention '%s' is not supported on TensorRT-RTX for Turing (SM 7.5). "
+            "Falling back to PyTorch.",
+            node.name,
+        )
+        return False
+
+    return True
+
+
 def scaled_dot_product_attention_validator(
     node: Node, settings: Optional[CompilationSettings] = None
 ) -> bool:
+    if not attention_capability_validator(node, settings):
+        return False
+
     attn_mask = args_bounds_check(node.args, 3, None)
     is_causal = args_bounds_check(node.args, 5, False)
     if is_causal and attn_mask is not None:
@@ -4639,6 +4677,9 @@ def aten_ops_scaled_dot_product_attention(
 def scaled_dot_product_flash_attention_validator(
     node: Node, settings: Optional[CompilationSettings] = None
 ) -> bool:
+    if not attention_capability_validator(node, settings):
+        return False
+
     if args_bounds_check(node.args, 5, False):
         _LOGGER.debug("return_debug_mask is not yet supported.")
         return False
@@ -4733,6 +4774,9 @@ def aten_ops_scaled_dot_product_flash_attention(
 def scaled_dot_product_efficient_attention_validator(
     node: Node, settings: Optional[CompilationSettings] = None
 ) -> bool:
+    if not attention_capability_validator(node, settings):
+        return False
+
     if args_bounds_check(node.args, 4, False):
         _LOGGER.debug("compute_log_sumexp is not yet supported.")
         return False
@@ -4819,6 +4863,7 @@ def scaled_dot_product_cudnn_attention_validator(
         _LOGGER.debug("return_debug_mask is not yet supported.")
         return False
 
+    # Delegating also picks up the Turing (SM 7.5) dtype guard.
     return scaled_dot_product_efficient_attention_validator(node, settings)
 
 

tests/py/dynamo/conversion/test_attention_aten.py

@@ -4,7 +4,7 @@
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 
-from .harness import DispatchTestCase
+from .harness import DispatchTestCase, skip_if_trt_rtx_turing
 
 
 class TestScaledDotProductAttention(DispatchTestCase):
@@ -112,6 +112,9 @@ def test_sdpa_bool_mask(
         dropout_p=0.0,
         enable_gqa=False,
     ):
+        if dtype == torch.float32:
+            skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention")
+
         class SDPA(nn.Module):
             def forward(self, query, key, value, attn_mask=None):
                 return torch.ops.aten.scaled_dot_product_attention.default(
@@ -247,6 +250,9 @@ def test_sdpa_fp_mask(
         dropout_p=0.0,
         enable_gqa=False,
     ):
+        if dtype == torch.float32:
+            skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention")
+
         class SDPA(nn.Module):
             def forward(self, query, key, value, attn_mask=None):
                 return torch.ops.aten.scaled_dot_product_attention.default(
@@ -352,6 +358,9 @@ def test_dynamic_sdpa_fp_mask(
         dropout_p=0.0,
         enable_gqa=False,
     ):
+        if dtype == torch.float32:
+            skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention")
+
         class SDPA(nn.Module):
             def forward(self, query, key, value, attn_mask=None):
                 return torch.ops.aten.scaled_dot_product_attention.default(
@@ -484,6 +493,9 @@ def test_efficient_sdpa(
         dtype,
         dropout_p=0.0,
     ):
+        if dtype == torch.float32:
+            skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention")
+
         class EfficientSDPA(nn.Module):
             def forward(self, query, key, value, attn_bias=None):
                 attn = torch.ops.aten._scaled_dot_product_efficient_attention.default(
@@ -606,6 +618,9 @@ def test_efficient_sdpa_random_attn_bias(
         dtype,
         dropout_p=0.0,
     ):
+        if dtype == torch.float32:
+            skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention")
+
         class EfficientSDPA(nn.Module):
             def forward(self, query, key, value, attn_bias=None):
                 attn = torch.ops.aten._scaled_dot_product_efficient_attention.default(

tests/py/dynamo/conversion/test_linear_aten.py

@@ -4,7 +4,7 @@
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 
-from .harness import DispatchTestCase
+from .harness import DispatchTestCase, skip_if_trt_rtx_turing
 
 
 class TestLinearConverter(DispatchTestCase):
@@ -17,6 +17,8 @@ class TestLinearConverter(DispatchTestCase):
         ]
     )
     def test_linear_converter(self, in_features, out_features):
+        skip_if_trt_rtx_turing(self, "nn.Linear (an FP32 GEMM)")
+
         class LinearModel(nn.Module):
             def __init__(self, in_features, out_features):
                 super(LinearModel, self).__init__()
@@ -30,6 +32,8 @@ def forward(self, x):
         self.run_test(model, inputs, use_dynamo_tracer=True, enable_passes=True)
 
     def test_linear_with_dynamic_shape(self):
+        skip_if_trt_rtx_turing(self, "aten.linear (an FP32 GEMM)")
+
         class LinearModel(torch.nn.Module):
             def forward(self, x, weight, bias):
                 return torch.ops.aten.linear.default(x, weight, bias)
@@ -50,6 +54,8 @@ def forward(self, x, weight, bias):
         )
 
     def test_linear_with_rank_3_input_and_bias(self):
+        skip_if_trt_rtx_turing(self, "aten.linear (an FP32 GEMM)")
+
         class LinearModel(torch.nn.Module):
             def forward(self, x, weight, bias):
                 return torch.ops.aten.linear.default(x, weight, bias)

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -7,6 +7,7 @@
 import torch.nn.functional as F
 import torch_tensorrt as torchtrt
 from torch import nn
+from torch_tensorrt._utils import trt_rtx_targets_turing
 from torch_tensorrt.dynamo._compiler import (
     convert_exported_program_to_serialized_trt_engine,
 )
@@ -19,6 +20,25 @@
 assertions = unittest.TestCase()
 
 
+def _no_fallback_gemm_dtype() -> torch.dtype:
+    """Model dtype for the tests that convert a whole program with no partitioner.
+
+    TensorRT-RTX cannot serve an FP32 GEMM on Turing (SM 7.5), so
+    ``gemm_capability_validator`` rejects the models' trailing ``nn.Linear`` there. In a
+    normal ``torchtrt.dynamo.compile`` the partitioner turns that rejection into a
+    PyTorch block, but ``convert_exported_program_to_serialized_trt_engine`` emits one
+    engine for the whole program by design and has nothing to fall back to, so the
+    rejection raises ``UnsupportedOperatorException`` instead.
+
+    These tests are about kwarg plumbing, not about FP32, so follow the capability
+    rather than switching them off: the guard keys on operand dtype only and Turing has
+    FP16 GEMM hardware. Keyed off the same predicate the validator uses so the two
+    cannot drift, and no ``enabled_precisions`` change is needed because the network is
+    strongly typed.
+    """
+    return torch.float16 if trt_rtx_targets_turing() else torch.float32
+
+
 @pytest.mark.unit
 @pytest.mark.critical
 def test_custom_model(tmpdir):
@@ -496,8 +516,9 @@ def forward(self, x, b=5, c=None, d=None):
                 x = x - d["value"]
             return self.fc1(x)
 
-    model = net().eval().to("cuda")
-    args = [torch.rand((1, 3, 224, 224)).to("cuda")]
+    dtype = _no_fallback_gemm_dtype()
+    model = net().eval().to("cuda").to(dtype)
+    args = [torch.rand((1, 3, 224, 224), dtype=dtype).to("cuda")]
     kwargs = {
         "b": torch.tensor(6).to("cuda"),
         "d": {"value": torch.tensor(8).to("cuda")},
@@ -1090,9 +1111,10 @@ def forward(self, x, b=5, c=None, d=None):
                 x = x - d["value"]
             return self.fc1(x)
 
-    model = net().eval().to("cuda")
+    dtype = _no_fallback_gemm_dtype()
+    model = net().eval().to("cuda").to(dtype)
     kwargs = {
-        "x": torch.rand((1, 3, 224, 224)).to("cuda"),
+        "x": torch.rand((1, 3, 224, 224), dtype=dtype).to("cuda"),
         "b": torch.tensor(6).to("cuda"),
         "d": {"value": torch.tensor(8).to("cuda")},
     }

tests/py/dynamo/models/test_model_refit.py

@@ -11,7 +11,7 @@
 import torch_tensorrt as torch_trt
 from torch import nn
 from torch_tensorrt._utils import is_tensorrt_rtx_version_supported
-from torch_tensorrt.dynamo import refit_module_weights
+from torch_tensorrt.dynamo import partitioning, refit_module_weights
 from torch_tensorrt.dynamo._refit import (
     construct_refit_mapping,
     get_engine_from_encoded_engine,
@@ -68,7 +68,26 @@ def test_mapping():
     )
     new_gm = exp_program2.module()
     new_gm = post_lowering(new_gm, settings)
-    mapping = construct_refit_mapping(new_gm, trt_input, settings)
+
+    # construct_refit_mapping interprets whatever module it is handed, so it has to be
+    # handed the same subgraph the engine above was built from. refit_module_weights
+    # always partitions first and maps each accelerated submodule; this test was the
+    # only caller passing a whole, un-partitioned module. That only worked because the
+    # model happened to be fully convertible: on Turing (SM 7.5) TensorRT-RTX cannot
+    # serve resnet18's FP32 fc GEMM, so capability partitioning routes it to PyTorch and
+    # the engine does not contain it -- while the un-partitioned module still does, with
+    # no partitioner in construct_refit_mapping to fall back to. Partition here the same
+    # way compile_module does, so the mapping matches the engine on either architecture.
+    num_supported_ops, total_ops = partitioning.get_graph_converter_support(
+        new_gm, settings.torch_executed_ops
+    )
+    partitioned_gm, _ = partitioning.fast_partition(
+        new_gm,
+        min_block_size=min_block_size,
+        torch_executed_ops=settings.torch_executed_ops,
+        skip_fusion=(num_supported_ops == total_ops),
+    )
+    mapping = construct_refit_mapping(partitioned_gm._run_on_acc_0, trt_input, settings)
 
     refitter = trt.Refitter(engine, TRT_LOGGER)
     weight_list = refitter.get_all_weights()