fix: fix automatic plugin test issue

Author: bowang007

.github/workflows/build-test-linux-x86_64.yml

@@ -426,7 +426,9 @@ jobs:
         set -euo pipefail
         pushd .
         cd tests/py/dynamo
-        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/l2_dynamo_plugins_tests_results.xml automatic_plugin/
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/l2_dynamo_plugins_tests_results.xml automatic_plugin/test_automatic_plugin.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/l2_dynamo_plugins_tests_results.xml automatic_plugin/test_automatic_plugin_with_attrs.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/l2_dynamo_plugins_tests_results.xml automatic_plugin/test_flashinfer_rmsnorm.py.py
         popd
 
   L2-torchscript-tests:

tests/py/dynamo/automatic_plugin/test_automatic_plugin.py

@@ -3,12 +3,13 @@
 
 import torch
 import torch.nn as nn
-import torch_tensorrt
 import triton
 import triton.language as tl
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 
+import torch_tensorrt
+
 from ..conversion.harness import DispatchTestCase
 
 
@@ -39,7 +40,7 @@ def elementwise_mul(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
     Z = torch.empty_like(X)
 
     # Define block size
-    BLOCK_SIZE = 1024
+    BLOCK_SIZE = 64
 
     # Grid of programs
     grid = lambda meta: (X.numel() // meta["BLOCK_SIZE"],)

tests/py/dynamo/automatic_plugin/test_automatic_plugin_with_attrs.py

@@ -3,12 +3,13 @@
 
 import torch
 import torch.nn as nn
-import torch_tensorrt
 import triton
 import triton.language as tl
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 
+import torch_tensorrt
+
 from ..conversion.harness import DispatchTestCase
 
 
@@ -40,7 +41,7 @@ def elementwise_scale_mul(
     Z = torch.empty_like(X)
 
     # Define block size
-    BLOCK_SIZE = 1024
+    BLOCK_SIZE = 64
 
     # Grid of programs
     grid = lambda meta: (X.numel() // meta["BLOCK_SIZE"],)

tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py

@@ -4,9 +4,10 @@
 import pytest
 import torch
 import torch.nn as nn
-import torch_tensorrt
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+
+import torch_tensorrt
 from torch_tensorrt._enums import dtype
 
 from ..conversion.harness import DispatchTestCase