diff --git a/.github/workflows/build-test-linux-x86_64.yml b/.github/workflows/build-test-linux-x86_64.yml index da3124c472..8bdb96ae96 100644 --- a/.github/workflows/build-test-linux-x86_64.yml +++ b/.github/workflows/build-test-linux-x86_64.yml @@ -426,7 +426,10 @@ 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}/dynamo_converters_test_results.xml -n 4 conversion/ + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_automatic_plugin.py + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_automatic_plugin_with_attrs.py + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_flashinfer_rmsnorm.py popd L2-torchscript-tests: diff --git a/tests/py/dynamo/automatic_plugin/test_automatic_plugin.py b/tests/py/dynamo/automatic_plugin/test_automatic_plugin.py index 44acb7b105..622b88547e 100644 --- a/tests/py/dynamo/automatic_plugin/test_automatic_plugin.py +++ b/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"],) diff --git a/tests/py/dynamo/automatic_plugin/test_automatic_plugin_with_attrs.py b/tests/py/dynamo/automatic_plugin/test_automatic_plugin_with_attrs.py index 823d0d600e..831a9f2d53 100644 --- a/tests/py/dynamo/automatic_plugin/test_automatic_plugin_with_attrs.py +++ b/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"],) diff --git a/tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py b/tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py index d85c8a633f..1666102070 100644 --- a/tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py +++ b/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