skip fp16 tests if TensorRT-RTX < 1.6

Author: zewenli98

py/torch_tensorrt/_utils.py

@@ -40,6 +40,22 @@ def check_cross_compile_trt_win_lib() -> bool:
     return False
 
 
+def is_tensorrt_rtx_version_supported(min_version: str) -> bool:
+    """
+    Check if the installed TensorRT-RTX version supports the specified minimum version.
+    Args:
+        min_version (str): Minimum required TensorRT-RTX version
+    Returns:
+        bool: True if TensorRT-RTX version is >= min_version, False otherwise
+    """
+    if trt._package_name != "tensorrt_rtx":
+        return True
+
+    from packaging.version import Version
+
+    return bool(Version(trt.__version__) >= Version(min_version))
+
+
 def is_tensorrt_version_supported(min_version: str) -> bool:
     """
     Check if the installed TensorRT version supports the specified minimum version.

tests/py/dynamo/models/test_hf_gqa_model.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 import torch_tensorrt
+from torch_tensorrt._utils import is_tensorrt_rtx_version_supported
 
 if importlib.util.find_spec("transformers"):
     import transformers
@@ -21,6 +22,15 @@
     "transformers is required to run this test",
 )
 def test_dynamic_head_dim_with_hf_model(dtype, decompose_attention):
+    if (
+        torch_tensorrt.ENABLED_FEATURES.tensorrt_rtx
+        and not is_tensorrt_rtx_version_supported("1.6")
+    ):
+        pytest.skip(
+            "TensorRT-RTX >= 1.6 is required because it fixed some accuracy issues for FP16"
+        )
+        return
+
     from transformers import AutoModelForCausalLM
 
     model_name = "Qwen/Qwen2.5-0.5B-Instruct"