Fix unstable CI (#299)

Author: jansel

test/test_autotuner.py

@@ -45,6 +45,7 @@ def test_config_fragment0(self):
         self.assertExpectedJournal("\n".join(map(repr, configs)))
 
     @patch.object(_compat, "_supports_tensor_descriptor", lambda: True)
+    @patch.object(loops, "_supports_warp_specialize", lambda: True)
     def test_config_fragment1(self):
         args = (
             torch.randn([8, 512, 512], device=DEVICE),

test/test_register_tunable.expected

@@ -2,7 +2,7 @@ This file is automatically generated by assertExpectedJournal calls in test_regi
 Update expected outputs by running tests with the EXPECTTEST_ACCEPT=1 environment variable set.
 
 --- assertExpectedJournal(TestRegisterTunable.test_integer_fragment)
-helion.Config(block_sizes=[128], range_unroll_factors=[0], range_warp_specializes=[None], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=4, num_stages=3, indexing='pointer', pid_type='flat', multiplier=3)
+helion.Config(block_sizes=[128], range_unroll_factors=[0], range_warp_specializes=[], range_num_stages=[0], range_multi_buffers=[None], range_flattens=[None], num_warps=4, num_stages=3, indexing='pointer', pid_type='flat', multiplier=3)
 
 --- assertExpectedJournal(TestRegisterTunable.test_integer_fragment)
 from __future__ import annotations

test/test_register_tunable.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
 
 import unittest
+from unittest.mock import patch
 
 import torch
 
 import helion
+from helion import _compat
 from helion._testing import DEVICE
 from helion._testing import TestCase
 from helion._testing import code_and_output
 from helion.autotuner import EnumFragment
 from helion.autotuner import IntegerFragment
 from helion.autotuner import PowerOfTwoFragment
 import helion.language as hl
+from helion.language import loops
 
 
 class TestRegisterTunable(TestCase):
@@ -41,6 +44,8 @@ def kernel_with_tunable(x: torch.Tensor) -> torch.Tensor:
         )
         self.assertExpectedJournal(code)
 
+    @patch.object(_compat, "_supports_tensor_descriptor", lambda: False)
+    @patch.object(loops, "_supports_warp_specialize", lambda: False)
     def test_integer_fragment(self):
         @helion.kernel()
         def kernel_with_int_param(x: torch.Tensor) -> torch.Tensor:

test/test_tensor_descriptor.expected

@@ -80,9 +80,7 @@ def attention(q_in: torch.Tensor, k_in: torch.Tensor, v_in: torch.Tensor):
     k_view = k_in.reshape([-1, n_dim, head_dim]).transpose(1, 2)
     out = torch.empty_like(q_view)
     sm_scale = 1.0 / math.sqrt(head_dim)
-    qk_scale = sm_scale * 1.44269504
     _BLOCK_SIZE_1 = 16
-    _RDIM_SIZE_2 = 64
     _BLOCK_SIZE_3 = 16
     _attention_kernel[q_in.size(1) * triton.cdiv(m_dim, _BLOCK_SIZE_1),](q_view, k_view, v_view, out, k_view.size(0), k_view.size(2), out.size(0), out.size(1), q_in.size(1), q_view.size(0), q_view.size(1), v_view.size(0), v_view.size(1), k_view.stride(0), k_view.stride(1), k_view.stride(2), out.stride(0), out.stride(1), out.stride(2), q_view.stride(0), q_view.stride(1), q_view.stride(2), v_view.stride(0), v_view.stride(1), v_view.stride(2), m_dim, n_dim, _BLOCK_SIZE_1, _BLOCK_SIZE_3, num_warps=4, num_stages=3)
     return out.view(q_in.size())
@@ -98,9 +96,7 @@ def _attention_make_precompiler(q_in: torch.Tensor, k_in: torch.Tensor, v_in: to
     k_view = k_in.reshape([-1, n_dim, head_dim]).transpose(1, 2)
     out = torch.empty_like(q_view)
     sm_scale = 1.0 / math.sqrt(head_dim)
-    qk_scale = sm_scale * 1.44269504
     _BLOCK_SIZE_1 = 16
-    _RDIM_SIZE_2 = 64
     _BLOCK_SIZE_3 = 16
     from helion.runtime.precompile_shim import make_precompiler
     return make_precompiler(_attention_kernel)(q_view, k_view, v_view, out, k_view.size(0), k_view.size(2), out.size(0), out.size(1), q_in.size(1), q_view.size(0), q_view.size(1), v_view.size(0), v_view.size(1), k_view.stride(0), k_view.stride(1), k_view.stride(2), out.stride(0), out.stride(1), out.stride(2), q_view.stride(0), q_view.stride(1), q_view.stride(2), v_view.stride(0), v_view.stride(1), v_view.stride(2), m_dim, n_dim, _BLOCK_SIZE_1, _BLOCK_SIZE_3, num_warps=4, num_stages=3)
@@ -182,9 +178,7 @@ def attention(q_in: torch.Tensor, k_in: torch.Tensor, v_in: torch.Tensor):
     k_view = k_in.reshape([-1, n_dim, head_dim]).transpose(1, 2)
     out = torch.empty_like(q_view)
     sm_scale = 1.0 / math.sqrt(head_dim)
-    qk_scale = sm_scale * 1.44269504
     _BLOCK_SIZE_1 = 128
-    _RDIM_SIZE_2 = 64
     _BLOCK_SIZE_3 = 64
     _attention_kernel[64 * triton.cdiv(1024, _BLOCK_SIZE_1),](q_view, k_view, v_view, out, _BLOCK_SIZE_1, _BLOCK_SIZE_3, num_warps=4, num_stages=3)
     return out.view(q_in.size())
@@ -200,9 +194,7 @@ def _attention_make_precompiler(q_in: torch.Tensor, k_in: torch.Tensor, v_in: to
     k_view = k_in.reshape([-1, n_dim, head_dim]).transpose(1, 2)
     out = torch.empty_like(q_view)
     sm_scale = 1.0 / math.sqrt(head_dim)
-    qk_scale = sm_scale * 1.44269504
     _BLOCK_SIZE_1 = 128
-    _RDIM_SIZE_2 = 64
     _BLOCK_SIZE_3 = 64
     from helion.runtime.precompile_shim import make_precompiler
     return make_precompiler(_attention_kernel)(q_view, k_view, v_view, out, _BLOCK_SIZE_1, _BLOCK_SIZE_3, num_warps=4, num_stages=3)