minor tests

nullplay · nullplay · commit f59fbd005fc3 · 2025-08-11T13:05:43.000-04:00
diff --git a/test/distributed/tensor/test_dtensor_compile.py b/test/distributed/tensor/test_dtensor_compile.py
@@ -891,7 +891,8 @@ def forward(self, input):
         FileCheck().check(
             "buf0 = torch.ops._c10d_functional.all_gather_into_tensor.default(primal"
         ).check("torch.ops._c10d_functional.wait_tensor.default(buf0").check(
-            "extern_kernels.mm(buf0,"
+            "extern_kernels.mm(buf0," if not config.triton.enable_native_matmul
+            else "triton_per_fused_add_addmm_0.run(buf6"
         ).run(code)
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -2718,6 +2718,9 @@ def forward(self, x, y):
                 result_package = model_package(*inputs_on_device)
             self.assertTrue(same(result_ref.cpu(), result_package.cpu()))
 
+    @unittest.skipIf(
+        config.triton.enable_native_matmul, "sin and mm are fused in native matmul"
+    )
     def test_reuse_kernel(self):
         class Model(torch.nn.Module):
             def __init__(self) -> None:
@@ -2736,14 +2739,9 @@ def forward(self, x, y):
         )
         model = Model()
 
-        if config.triton.enable_native_matmul:
-            atol, rtol = 2e-4, 2e-4
-        else :
-            # 1e-4 is the tol value used in pytorch/torch/_dynamo/utils.py
-            atol, rtol = 1e-4, 1e-4
-
+        # 1e-4 is the tol value used in pytorch/torch/_dynamo/utils.py
         self.check_model(
-            model, example_inputs, atol=atol, rtol=rtol
+            model, example_inputs, atol=1e-4, rtol=1e-4
         )  
 
         if self.device == "mps":
@@ -4921,7 +4919,10 @@ def forward(self, image: torch.Tensor, target_size: torch.Tensor):
             "target_size": None,
         }
         self.check_model(model, example_inputs, dynamic_shapes=dynamic_shapes)
-
+    
+    @unittest.skipIf(
+        config.triton.enable_native_matmul, "matmul is generated"
+    )
     def test_aoti_debug_printer_codegen(self):
         # basic addmm model to test codegen for aoti intermediate debug printer
         class Model(torch.nn.Module):
diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py
@@ -3650,6 +3650,9 @@ def mask_mod(b, h, q, kv):
 
     @supported_platform
     @skip_on_cpu
+    @unittest.skipIf(
+        config.triton.enable_native_matmul, "different dynamo counters"
+    )
     def test_free_symbol_dynamic(self, device):
         def batch_flip_causal(b, h, q_idx, kv_idx):
             return (q_idx >= kv_idx) & (b % 2 == 0)
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -6304,6 +6304,9 @@ def fn(x1, x2, x3, x4):
     @skip_if_gpu_halide
     # Constant folding was explicitly turned off due to issue #108388
     # Turn it back on for test
+    @unittest.skipIf(
+        config.triton.enable_native_matmul, "native matmul has better precision"
+    )
     @torch._inductor.config.patch(joint_graph_constant_folding=True)
     def test_remove_no_ops(self):
         def matmul_with_op(x, y, fn):
@@ -6329,11 +6332,7 @@ def matmul_with_op(x, y, fn):
             if self.device == "cpu":
                 FileCheck().check_not("cpp_fused").run(source_codes[0])
             else :
-                if config.triton.enable_native_matmul: 
-                    FileCheck().check("triton.jit").run(source_codes[0])
-                    #atol, rtol = 1e-2, 1e-2
-                else :
-                    FileCheck().check_not("triton.jit").run(source_codes[0])
+                FileCheck().check_not("triton.jit").run(source_codes[0])
 
         # test dtype conversion
         for lowp_dtype in [torch.float16, torch.bfloat16]:
