chore: add moe tests

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

@@ -151,7 +151,6 @@ jobs:
                 python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_runtime_tests_results.xml runtime/test_000_*
                 python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_partitioning_tests_results.xml partitioning/test_000_*
                 python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_core_lowering_tests_results.xml lowering/
-                python -m pytest -ra -n 8 --junitxml=${RUNNER_TEST_RESULTS_DIR}/l0_dynamo_hlo_tests_results.xml hlo/
                 popd
 
     L0-py-core-tests:

tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py

@@ -59,7 +59,7 @@ def forward(self, input, weight):
             torch.randn(weight_shape, device="cuda", dtype=data_type),
         ]
 
-        self.run_test(rmsnorm(), inputs, precision=dtype.f16)
+        self.run_test(rmsnorm(), inputs)
 
 
 if __name__ == "__main__":

tests/py/dynamo/conversion/harness.py

@@ -272,7 +272,7 @@ def run_test_with_error(self, mod, inputs, interpreter, expect_error):
                     cuda_inputs.append(i.cuda())
 
                 mod.eval()
-                interpreter.run(precision=torch.float)
+                interpreter.run()
 
     def assert_has_op(self, mod, ops):
         ops_in_mod = set()
@@ -359,7 +359,6 @@ def run_test(
         inputs,
         rtol=RTOL,
         atol=ATOL,
-        precision=dtype.f32,
         check_dtype=True,
         use_dynamo_tracer=None,
         enable_passes=False,
@@ -368,6 +367,8 @@ def run_test(
         immutable_weights=True,
         decompose_attention=False,
         attn_bias_is_causal=True,
+        require_full_compilation=False,
+        disable_tf32=False,
     ):
         # TODO: lan to remove this and set use_dynamo_traccer to True by default
         # once all the converter test files are moved to use_dynamo_tracer
@@ -379,6 +380,8 @@ def run_test(
             immutable_weights=immutable_weights,
             decompose_attention=decompose_attention,
             attn_bias_is_causal=attn_bias_is_causal,
+            require_full_compilation=require_full_compilation,
+            disable_tf32=disable_tf32,
         )
 
         mod = self.generate_graph(
@@ -444,6 +447,13 @@ def run_test(
             compilation_settings=compilation_settings,
         )
 
+        if require_full_compilation:
+            missing = interp.validate_conversion()
+            self.assertTrue(
+                len(missing) == 0,
+                f"require_full_compilation=True but the following ops don't have TRT converter: {missing}",
+            )
+
         super().run_test(
             mod,
             trt_inputs,
@@ -460,7 +470,6 @@ def run_test_compare_tensor_attributes_only(
         inputs,
         expected_ops,
         comparators: List[Tuple[Callable, List]],
-        precision=torch.float,
         output_dtypes=None,
         use_dynamo_tracer=False,
         enable_passes=False,

tests/py/dynamo/conversion/test_attention.py

@@ -27,7 +27,6 @@ def forward(self, query, key, value):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=torch.float16,
             enable_passes=True,
             decompose_attention=True,
         )
@@ -63,7 +62,6 @@ def forward(self, query, key, value):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=torch.float16,
             enable_passes=True,
             decompose_attention=True,
         )
@@ -96,7 +94,6 @@ def forward(self, query, key, value):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=torch.float16,
             enable_passes=True,
             decompose_attention=True,
         )

tests/py/dynamo/conversion/test_attention_aten.py

@@ -139,7 +139,6 @@ def forward(self, query, key, value, attn_mask=None):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=dtype,
             enable_passes=True,
         )
 
@@ -274,7 +273,6 @@ def forward(self, query, key, value, attn_mask=None):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=dtype,
             enable_passes=True,
         )
 
@@ -513,7 +511,6 @@ def forward(self, query, key, value, attn_bias=None):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=dtype,
             enable_passes=True,
         )
 
@@ -630,7 +627,6 @@ def forward(self, query, key, value, attn_bias=None):
             inputs,
             rtol=1e-2,
             atol=1e-2,
-            precision=dtype,
             enable_passes=True,
         )
 

tests/py/dynamo/conversion/test_casts.py

@@ -47,11 +47,10 @@ def forward(self, x):
                 y = torch.ops.aten._to_copy.default(x, dtype=torch.half)
                 return y
 
-        inputs = [torch.rand((1, 3, 10))]
+        inputs = [torch.rand((1, 3, 10), dtype=torch.half)]
         self.run_test(
             ToCopyHalf(),
             inputs,
-            precision=torch.half,
         )
 
     def test_to_copy_float(self):
@@ -60,11 +59,10 @@ def forward(self, x):
                 y = torch.ops.aten._to_copy.default(x, dtype=torch.float)
                 return y
 
-        inputs = [torch.rand((1, 3, 10)).half()]
+        inputs = [torch.rand((1, 3, 10), dtype=torch.float)]
         self.run_test(
             ToCopyFloat(),
             inputs,
-            precision=torch.float,
         )
 
     def test_to_copy_bfloat16(self):
@@ -74,11 +72,10 @@ def forward(self, x):
                 y = y**2
                 return y
 
-        inputs = [torch.rand((1, 3, 10), dtype=torch.float32)]
+        inputs = [torch.rand((1, 3, 10), dtype=torch.bfloat16)]
         self.run_test(
             ToCopyBFloat16(),
             inputs,
-            precision=torch.float,
         )
 
     def test_to_copy_i64b(self):
@@ -102,11 +99,10 @@ def forward(self, x):
                 z = torch.ops.aten._to_copy.default(x_1, dtype=torch.float)
                 return y, z
 
-        inputs = [torch.rand((1, 3, 10))]
+        inputs = [torch.rand((1, 3, 10), dtype=torch.float)]
         self.run_test(
             ToCopyReturns(),
             inputs,
-            precision=torch.float,
         )
 
 

tests/py/dynamo/conversion/test_embedding_bag_aten.py

@@ -146,7 +146,6 @@ def forward(self, weight, indices):
         self.run_test(
             TestEmbeddingBag(),
             inputs=[weight, indices],
-            precision=weight.dtype,
             enable_passes=True,
             propagate_shapes=True,
             immutable_weights=True,
@@ -345,7 +344,6 @@ def forward(self, weight, indices, offsets):
         self.run_test(
             TestEmbeddingBag(),
             inputs=[weight, indices, offsets],
-            precision=weight.dtype,
             enable_passes=True,
             propagate_shapes=True,
             immutable_weights=True,
@@ -410,7 +408,6 @@ def forward(self, weight, indices, offsets):
         self.run_test(
             TestEmbeddingBag(),
             inputs=[weight, indices, offsets],
-            precision=weight.dtype,
             enable_passes=True,
             propagate_shapes=True,
             immutable_weights=True,

tests/py/dynamo/conversion/test_erf_aten.py

@@ -22,7 +22,7 @@ def forward(self, input):
                 return torch.ops.aten.erf.default(input)
 
         inputs = [torch.randn(x, dtype=type)]
-        self.run_test(erf(), inputs, precision=type)
+        self.run_test(erf(), inputs)
 
     @parameterized.expand(
         [

tests/py/dynamo/conversion/test_group_norm_aten.py

@@ -42,7 +42,6 @@ def forward(self, x, weight, bias):
         self.run_test(
             GroupNorm(),
             inputs,
-            precision=torch.half,
             use_dynamo_tracer=True,
             enable_passes=True,
         )

tests/py/dynamo/conversion/test_hard_sigmoid_aten.py

@@ -51,11 +51,10 @@ class TestModule(nn.Module):
             def forward(self, x):
                 return torch.ops.aten.hardsigmoid.default(x)
 
-        inputs = [torch.randn(1, 10)]
+        inputs = [torch.randn(1, 10, dtype=torch.float16)]
         self.run_test(
             TestModule(),
             inputs,
-            precision=torch.half,
             check_dtype=False,
         )
 

tests/py/dynamo/conversion/test_neg_aten.py

@@ -21,8 +21,8 @@ class neg(nn.Module):
             def forward(self, input):
                 return torch.ops.aten.neg.default(input)
 
-        inputs = [torch.randn(x, dtype=type)]
-        self.run_test(neg(), inputs, precision=type)
+        inputs = [torch.randn(x, dtype=type).cuda()]
+        self.run_test(neg(), inputs)
 
     @parameterized.expand(
         [

tests/py/dynamo/conversion/test_sigmoid_aten.py

@@ -49,11 +49,10 @@ class TestModule(nn.Module):
             def forward(self, x):
                 return torch.ops.aten.sigmoid.default(x)
 
-        inputs = [torch.randn(1, 10)]
+        inputs = [torch.randn(1, 10, dtype=torch.float16)]
         self.run_test(
             TestModule(),
             inputs,
-            precision=torch.half,
             check_dtype=False,
         )