Repository: pytorch/TensorRT (language: python)

Issue:
🐛 [Bug] TS test_scaled_dot_product_attention failed
In TRT 10.3 upgrade, TorchScript `tests/core/conversion/converters/test_scaled_dot_product_attention.cpp` failed due to the error:

[  FAILED  ] Converters.ATenScaledDotProductAttentionConvertsCorrectly
[  FAILED  ] Converters.ATenScaledDotProductAttnMaskFloatConvertsCorrectly
[  FAILED  ] Converters.ATenScaledDotProductAttnMaskBoolConvertsCorrectly
	unknown file: Failure                                                                                                                                    C++ exception with description "0 INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":617, please report a bug to PyTorch. We don't have an op for aten::scaled_dot_product_attention but it isn't a special case.  Argument types: Tensor, Tensor, Tensor, Tensor, float, bool, NoneType,                                                                                                                                                                 Candidates:                                                                                                                                                      aten::scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0., bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> Tensor
Exception raised from analyzeImpl at ../torch/csrc/jit/ir/alias_analysis.cpp:617 (most recent call first):                                               frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x96 (0x7f543dc12c16 in /opt/python3/site-packages/torch/lib/libc10.so)                  frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7f543dbbf462 in /opt/python3/site-packages/torch/lib/libc10.so)                                                                                                                                       frame #2: c10::detail::torchInternalAssertFail(char const*, char const*, unsigned int, char const*, std::string const&) + 0x3f (0x7f543dc0fb0f in /opt/python3/site-packages/torch/lib/libc10.so)                                                                                                                 frame #3: <unknown function> + 0x546ec98 (0x7f548add8c98 in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)
frame #4: <unknown function> + 0x5470027 (0x7f548adda027 in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)                                        frame #5: torch::jit::AliasDb::AliasDb(std::shared_ptr<torch::jit::Graph>, bool, bool) + 0x267 (0x7f548addb5e7 in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)                                                                                                                                           frame #6: <unknown function> + 0x5555fad (0x7f548aebffad in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)                                        frame #7: torch::jit::EliminateDeadCode(std::shared_ptr<torch::jit::Graph> const&, torch::jit::DCESideEffectPolicy) + 0x1e5 (0x7f548aeb45e5 in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)
frame #8: torch::jit::CanonicalizeOps(std::shared_ptr<torch::jit::Graph> const&) + 0x32d (0x7f548ae8ddcd in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)
frame #9: torch::jit::ProfilingGraphExecutorImpl::runProfilingInsensitiveOptimizations(std::shared_ptr<torch::jit::Graph>&) + 0x3f0 (0x7f548b0b5560 in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)
frame #10: torch::jit::ProfilingGraphExecutorImpl::getOptimizedPlanFor(std::vector<c10::IValue, std::allocator<c10::IValue> >&, std::optional<unsigned long>) + 0xadf (0x7f548b0bc59f in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)
frame #11: torch::jit::ProfilingGraphExecutorImpl::getPlanFor(std::vector<c10::IValue, std::allocator<c10::IValue> >&, std::optional<unsigned long>) + 0xcc (0x7f548b0bcb6c in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)                                                                              frame #12: <unknown function> + 0x570acc8 (0x7f548b074cc8 in /opt/python3/site-packages/torch/lib/libtorch_cpu.so)                                       
frame #13: torch_tensorrt::tests::util::RunGraph(std::shared_ptr<torch::jit::Graph>&, std::map<torch::jit::Value*, c10::IValue, std::less<torch::jit::Value*>, std::allocator<std::pair<torch::jit::Value* const, c10::IValue> > >&, std::vector<at::Tensor, std::allocator<at::Tensor> >) + 0x3e6 (0x7f54ac36c440 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libtests_Sutil_Slibutil.so)
frame #14: <unknown function> + 0x7c799 (0x556cc6c77799 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/sandbox/processwrapper-sandbox/486/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/test_scaled_dot_product_attention.runfiles/_main/tests/core/conversion/converters/test_scaled_dot_product_attention)
frame #15: void testing::internal::HandleSehExceptionsInMethodIfSupported<testing::Test, void>(testing::Test*, void (testing::Test::*)(), char const*) + 0x69 (0x7f5499facd13 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #16: void testing::internal::HandleExceptionsInMethodIfSupported<testing::Test, void>(testing::Test*, void (testing::Test::*)(), char const*) + 0x4f (0x7f5499fa71ae in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #17: testing::Test::Run() + 0x104 (0x7f5499f8990c in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #18: testing::TestInfo::Run() + 0x16d (0x7f5499f8a2f5 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #19: testing::TestSuite::Run() + 0x1bf (0x7f5499f8ac33 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #20: testing::internal::UnitTestImpl::RunAllTests() + 0x39f (0x7f5499f9977b in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #21: bool testing::internal::HandleSehExceptionsInMethodIfSupported<testing::internal::UnitTestImpl, bool>(testing::internal::UnitTestImpl*, bool (testing::internal::UnitTestImpl::*)(), char const*) + 0x69 (0x7f5499fadca1 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #22: bool testing::internal::HandleExceptionsInMethodIfSupported<testing::internal::UnitTestImpl, bool>(testing::internal::UnitTestImpl*, bool (testing::internal::UnitTestImpl::*)(), char const*) + 0x4f (0x7f5499fa8304 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #23: testing::UnitTest::Run() + 0xbe (0x7f5499f97fdc in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest.so)
frame #24: RUN_ALL_TESTS() + 0x15 (0x7f549a3d7b24 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest_Umain.so)
frame #25: main + 0x44 (0x7f549a3d7a9d in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/../../../../_solib_k8/libexternal_Sgoogletest~_Slibgtest_Umain.so)
frame #26: <unknown function> + 0x29d90 (0x7f543d75cd90 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #27: __libc_start_main + 0x80 (0x7f543d75ce40 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #28: <unknown function> + 0x7ae15 (0x556cc6c75e15 in /root/.cache/bazel/_bazel_root/208025f0ef0faeb3920e1c325d843c43/sandbox/processwrapper-sandbox/486/execroot/_main/bazel-out/k8-dbg/bin/tests/core/conversion/converters/test_scaled_dot_product_attention.runfiles/_main/tests/core/conversion/converters/test_scaled_dot_product_attention)
" thrown in the test body.


## To Reproduce

Steps to reproduce the behavior:
```
bazel test //tests/core/conversion/converters:test_scaled_dot_product_attention --compilation_mode=dbg --test_output=all --test_timeout=8000 --jobs 4 --config=pre_cxx11_abi --cxxopt='-std=c++17'
```

## Additional context

The error is probably because [the schema](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml#L14723) changed for the op (`bool enable_gqa=False`). We might have to update the signature in our lowering/converter.
fix: Relax thresholds of dynamo converters' tests
# Description

After bumping to TRT 10, some Dynamo converters' tests in Torch-TRT failed due to threshold/tolerance issue. TRT team suggests using 5e-3 as the threshold.

Fixes #3053 

## Type of change

Please delete options that are not relevant and/or add your own.

- Bug fix (non-breaking change which fixes an issue)

# Checklist:

- [ ] My code follows the style guidelines of this project (You can use the linters)
- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas and hacks
- [ ] I have made corresponding changes to the documentation
- [ ] I have added tests to verify my fix or my feature
- [ ] New and existing unit tests pass locally with my changes
- [ ] I have added the relevant labels to my PR in so that relevant reviewers are notified



Discussion / hints:

Change the TS tests as well

Fix patch:
diff --git a/core/lowering/passes/unpack_scaled_dot_product_attention.cpp b/core/lowering/passes/unpack_scaled_dot_product_attention.cpp
index 3c347f65ca..7d9f744d3f 100644
--- a/core/lowering/passes/unpack_scaled_dot_product_attention.cpp
+++ b/core/lowering/passes/unpack_scaled_dot_product_attention.cpp
@@ -12,12 +12,12 @@ namespace passes {
 // https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
 void UnpackScaledDotProductAttention(std::shared_ptr<torch::jit::Graph>& graph) {
   std::string sdpa_pattern = R"IR(
-    graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale):
-      %out: Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale)
+    graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale, %enable_gqa):
+      %out: Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale, %enable_gqa)
       return (%out))IR";
 
   std::string unpacked_sdpa_pattern = R"IR(
-    graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale):
+    graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale, %enable_gqa):
       %none : NoneType = prim::Constant()
       %1 : int = prim::Constant[value=-1]()
       %2 : int = prim::Constant[value=-2]()
@@ -33,7 +33,7 @@ void UnpackScaledDotProductAttention(std::shared_ptr<torch::jit::Graph>& graph)
       return(%out))IR";
 
   std::string unpacked_sdpa_attn_biased_pattern = R"IR(
-    graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale):
+    graph(%query, %key, %value, %attn_mask, %dropout_p, %is_causal, %scale, %enable_gqa):
       %none : NoneType = prim::Constant()
       %0 : int = prim::Constant[value=1]()
       %1 : int = prim::Constant[value=-1]()
@@ -69,6 +69,16 @@ void UnpackScaledDotProductAttention(std::shared_ptr<torch::jit::Graph>& graph)
         if (attn_mask_node->kind() != at::prim::Constant || !attn_mask_node->mustBeNone()) {
           return false;
         }
+        auto enable_gqa_node = match.anchor->inputs().at(7)->node();
+        if (enable_gqa_node->kind() != at::prim::Constant) {
+          LOG_WARNING(
+              "Could not unpack scaled_dot_product_attention with non constant enable_gqa: " << *enable_gqa_node);
+          return false;
+        }
+        if (enable_gqa_node->i(at::attr::value) == 1) {
+          LOG_WARNING("Could not unpack scaled_dot_product_attention with enable_gqa = True: " << *enable_gqa_node);
+          return false;
+        }
         return true;
       });
 
@@ -83,6 +93,11 @@ void UnpackScaledDotProductAttention(std::shared_ptr<torch::jit::Graph>& graph)
           // messages already written in first pass, do not write again
           return false;
         }
+        auto enable_gqa_node = match.anchor->inputs().at(7)->node();
+        if (enable_gqa_node->kind() != at::prim::Constant || enable_gqa_node->i(at::attr::value) == 1) {
+          // messages already written in first pass, do not write again
+          return false;
+        }
         return true;
       });
   LOG_GRAPH("Post unpack scaled_dot_product_attention: " << *graph);


Test patch:
diff --git a/tests/core/conversion/converters/test_scaled_dot_product_attention.cpp b/tests/core/conversion/converters/test_scaled_dot_product_attention.cpp
index 6809472529..0b62b57cee 100644
--- a/tests/core/conversion/converters/test_scaled_dot_product_attention.cpp
+++ b/tests/core/conversion/converters/test_scaled_dot_product_attention.cpp
@@ -11,8 +11,9 @@ TEST(Converters, ATenScaledDotProductAttentionConvertsCorrectly) {
         %none : NoneType = prim::Constant()
         %0 : float = prim::Constant[value=0.]()
         %scale : NoneType = prim::Constant()
+        %enable_gqa : bool = prim::Constant[value=0]()
         %false : bool = prim::Constant[value=0]()
-        %3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %none, %0, %false, %scale)
+        %3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %none, %0, %false, %scale, %enable_gqa)
         return (%3))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();
@@ -38,7 +39,8 @@ TEST(Converters, ATenScaledDotProductAttnMaskFloatConvertsCorrectly) {
         %0 : float = prim::Constant[value=0.]()
         %false : bool = prim::Constant[value=0]()
         %scale : NoneType = prim::Constant()
-        %3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %0, %false, %scale)
+        %enable_gqa : bool = prim::Constant[value=0]()
+        %3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %0, %false, %scale, %enable_gqa)
         return (%3))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();
@@ -59,13 +61,14 @@ TEST(Converters, ATenScaledDotProductAttnMaskFloatConvertsCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0]));
 }
 
-TEST(Converters, ATenScaledDotProductAttnMaskBoolConvertsCorrectly) {
+TEST(Converters, ATenScaledDotProductAttnMaskIntConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%query : Tensor, %key : Tensor, %value : Tensor, %attn_mask : Tensor):
         %0 : float = prim::Constant[value=0.]()
         %false : bool = prim::Constant[value=0]()
         %scale : NoneType = prim::Constant()
-        %3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %0, %false, %scale)
+        %enable_gqa : bool = prim::Constant[value=0]()
+        %3 : Tensor = aten::scaled_dot_product_attention(%query, %key, %value, %attn_mask, %0, %false, %scale, %enable_gqa)
         return (%3))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();
@@ -74,7 +77,7 @@ TEST(Converters, ATenScaledDotProductAttnMaskBoolConvertsCorrectly) {
   auto query = at::rand({32, 8, 128, 64}, {at::kCUDA});
   auto key = at::rand({32, 8, 128, 64}, {at::kCUDA});
   auto value = at::rand({32, 8, 128, 64}, {at::kCUDA});
-  auto attn_mask = at::randint(0, 2, {32, 8, 128, 128}, at::kCUDA).to(at::kBool);
+  auto attn_mask = at::randint(0, 2, {32, 8, 128, 128}, {at::kCUDA});
   auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
   auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {query, key, value, attn_mask});