More fair comparison

stack-info: PR: #146, branch: drisspg/stack/7

Author: drisspg

examples/attention.py

@@ -87,9 +87,10 @@ def test(
     p = torch.softmax(p.float(), dim=-1).to(dtype)
     ref_out = torch.matmul(p, v)
 
-    # flex attention version
+    # flex attention version=
     # TODO(jansel): turn the above kernel into a flex attention kernel
-    flex_out = flex_attention(q, k, v)
+    flex_compiled = torch.compile(flex_attention, fullgraph=True)
+    flex_out = flex_compiled(q, k, v)
     torch.testing.assert_close(flex_out, ref_out, atol=1e-2, rtol=1e-2)
 
     # sdpa version
@@ -106,7 +107,7 @@ def test(
     spda_sec = do_bench(
         lambda: torch.nn.functional.scaled_dot_product_attention(q, k, v)
     )
-    flex_sec = do_bench(lambda: flex_attention(q, k, v))
+    flex_sec = do_bench(lambda: flex_compiled(q, k, v))
     helion_sec = do_bench(lambda: attention(q, k, v))
     print(
         f"Helion time: {helion_sec:.4f}ms, flex time: {flex_sec:.4f}, torch time: {spda_sec:.4f}"