Fix test_fp8_attention expected (#355)

This was triggering autotuning and making tests take 10+ minutes

Author: jansel

helion/_testing.py

@@ -137,6 +137,8 @@ def check_example(
     fn_name: str | None = None,
     skip_accuracy: bool = False,
     static_shapes: bool | None = None,
+    atol: float = 1e-1,
+    rtol: float = 1e-2,
     **kwargs: object,
 ) -> str:
     """Helper used in unit tests to run a single example kernel and check its output."""
@@ -157,8 +159,8 @@ def check_example(
         torch.testing.assert_close(
             result.to(torch.float32),
             expected.to(torch.float32),
-            atol=1e-1,
-            rtol=1e-2,
+            atol=atol,
+            rtol=rtol,
         )
     return code
 

test/test_examples.py

@@ -587,7 +587,7 @@ def test_fp8_attention(self):
         args = (q_fp8, k_fp8, v_fp8, batch, heads)
 
         # Get expected output from kernel
-        expected = mod.fp8_attention_kernel(*args)
+        expected = mod.fp8_attention_pytorch(q, k, v)()
 
         self.assertExpectedJournal(
             check_example(
@@ -596,6 +596,8 @@ def test_fp8_attention(self):
                 expected,
                 fn_name="fp8_attention_kernel",
                 block_sizes=[64, 64],
+                atol=0.2,
+                rtol=0.1,
             )
         )