[kernels][tests] Use fresh_knobs when touching triton.knobs (#7687)

Follow up to #7682

Author: peterbell10

python/test/conftest.py

@@ -1,7 +1,5 @@
-import os
 import pytest
 import tempfile
-from typing import Optional, Set
 
 
 def pytest_configure(config):
@@ -26,49 +24,9 @@ def fresh_triton_cache():
             yield tmpdir
 
 
-def _fresh_knobs_impl(monkeypatch, skipped_attr: Optional[Set[str]] = None):
-    from triton import knobs
-
-    if skipped_attr is None:
-        skipped_attr = set()
-
-    knobs_map = {
-        name: knobset
-        for name, knobset in knobs.__dict__.items()
-        if isinstance(knobset, knobs.base_knobs) and knobset != knobs.base_knobs and name not in skipped_attr
-    }
-
-    # We store which variables we need to unset below in finally because
-    # monkeypatch doesn't appear to reset variables that were never set
-    # before the monkeypatch.delenv call below.
-    env_to_unset = []
-    prev_propagate_env = knobs.propagate_env
-
-    def fresh_function():
-        nonlocal env_to_unset
-        for name, knobset in knobs_map.items():
-            setattr(knobs, name, knobset.copy().reset())
-            for knob in knobset.knob_descriptors.values():
-                if knob.key in os.environ:
-                    monkeypatch.delenv(knob.key, raising=False)
-                else:
-                    env_to_unset.append(knob.key)
-        knobs.propagate_env = True
-        return knobs
-
-    def reset_function():
-        for name, knobset in knobs_map.items():
-            setattr(knobs, name, knobset)
-        for k in env_to_unset:
-            if k in os.environ:
-                del os.environ[k]
-        knobs.propagate_env = prev_propagate_env
-
-    return fresh_function, reset_function
-
-
 @pytest.fixture
 def fresh_knobs(monkeypatch):
+    from triton._internal_testing import _fresh_knobs_impl
     fresh_function, reset_function = _fresh_knobs_impl(monkeypatch)
     try:
         yield fresh_function()
@@ -83,6 +41,7 @@ def fresh_knobs_except_libraries(monkeypatch):
     information from the environment as these may be
     needed to successfully compile kernels.
     """
+    from triton._internal_testing import _fresh_knobs_impl
     fresh_function, reset_function = _fresh_knobs_impl(monkeypatch, skipped_attr={"build", "nvidia", "amd"})
     try:
         yield fresh_function()

python/triton/_internal_testing.py

@@ -5,10 +5,10 @@
 import triton
 import triton.language as tl
 from triton import knobs
+from typing import Optional, Set, Union
 import pytest
 
 from numpy.random import RandomState
-from typing import Optional, Union
 from triton.runtime.jit import TensorWrapper, reinterpret, type_canonicalisation_dict
 
 int_dtypes = ['int8', 'int16', 'int32', 'int64']
@@ -202,3 +202,44 @@ def unwrap_tensor(t: Union[torch.Tensor, triton.runtime.jit.TensorWrapper]) -> t
     if isinstance(t, triton.runtime.jit.TensorWrapper):
         return t.base
     return t
+
+
+def _fresh_knobs_impl(monkeypatch, skipped_attr: Optional[Set[str]] = None):
+    from triton import knobs
+
+    if skipped_attr is None:
+        skipped_attr = set()
+
+    knobs_map = {
+        name: knobset
+        for name, knobset in knobs.__dict__.items()
+        if isinstance(knobset, knobs.base_knobs) and knobset != knobs.base_knobs and name not in skipped_attr
+    }
+
+    # We store which variables we need to unset below in finally because
+    # monkeypatch doesn't appear to reset variables that were never set
+    # before the monkeypatch.delenv call below.
+    env_to_unset = []
+    prev_propagate_env = knobs.propagate_env
+
+    def fresh_function():
+        nonlocal env_to_unset
+        for name, knobset in knobs_map.items():
+            setattr(knobs, name, knobset.copy().reset())
+            for knob in knobset.knob_descriptors.values():
+                if knob.key in os.environ:
+                    monkeypatch.delenv(knob.key, raising=False)
+                else:
+                    env_to_unset.append(knob.key)
+        knobs.propagate_env = True
+        return knobs
+
+    def reset_function():
+        for name, knobset in knobs_map.items():
+            setattr(knobs, name, knobset)
+        for k in env_to_unset:
+            if k in os.environ:
+                del os.environ[k]
+        knobs.propagate_env = prev_propagate_env
+
+    return fresh_function, reset_function

python/triton_kernels/tests/conftest.py

@@ -8,3 +8,13 @@ def pytest_addoption(parser):
 @pytest.fixture
 def device(request):
     return request.config.getoption("--device")
+
+
+@pytest.fixture
+def fresh_knobs(monkeypatch):
+    from triton._internal_testing import _fresh_knobs_impl
+    fresh_function, reset_function = _fresh_knobs_impl(monkeypatch)
+    try:
+        yield fresh_function()
+    finally:
+        reset_function()

python/triton_kernels/tests/test_matmul.py

@@ -255,7 +255,7 @@ class Case:
 @pytest.mark.parametrize("is_persistent", [False, True])
 def test_op(m, n, k, split_k, do_gather, do_scatter, fused_scatter, has_y_gammas, is_persistent, n_expts_tot,
             n_expts_act, n_expt_shards, mode, act_dtype_str, weight_dtype_str, block_m, hbm_swizzling, epilogue_subtile,
-            device, opt_flags_scope):
+            device, opt_flags_scope, fresh_knobs):
     # TODO: remove when Triton FP8 supports proper RTNE
     if is_cuda():
         if "float8" in weight_dtype_str and torch.cuda.get_device_capability()[0] < 9:

python/triton_kernels/tests/test_specialize.py

@@ -38,7 +38,7 @@ def cacheable_kernel():
     return get_specialized_kernel()
 
 
-def test_cacheable(device):
+def test_cacheable(device, fresh_knobs):
     specialized_kernel = get_specialized_kernel()
 
     specialization_data = None
@@ -53,36 +53,32 @@ def cache_hook(*args, **kwargs):
         fn_name = kwargs["fn"].name
         module_name = kwargs["fn"].module
 
-    prev_hook = triton.knobs.runtime.jit_cache_hook
-    try:
-        triton.knobs.runtime.jit_cache_hook = cache_hook
-        o = torch.empty((1, ), dtype=torch.float32, device=device)
-        k = specialized_kernel[(1, )](o, )
-        hash = k.hash
-        assert o.item() == 1.0
-        assert module_name == "tests.test_specialize"
-        assert fn_name == "cacheable_kernel"
-
-        compile_count = 0
-
-        def count_hook(*args, **kwargs):
-            nonlocal compile_count
-            compile_count += 1
-
-        triton.knobs.runtime.jit_cache_hook = count_hook
-        # clear the cache
-        specialized_kernel.device_caches.clear()
-
-        # retrieve the kernel from name and preload it.
-        fn = retrieve_fn(module_name, fn_name)
-        assert fn == specialized_kernel
-        preload = fn.preload(specialization_data)
-        assert compile_count == 1
-        assert preload.hash == hash
-
-        # verify that we hit the cache.
-        compile_count = 0
-        specialized_kernel[(1, )](o, )
-        assert compile_count == 0
-    finally:
-        triton.knobs.runtime.jit_cache_hook = prev_hook
+    triton.knobs.runtime.jit_cache_hook = cache_hook
+    o = torch.empty((1, ), dtype=torch.float32, device=device)
+    k = specialized_kernel[(1, )](o, )
+    hash = k.hash
+    assert o.item() == 1.0
+    assert module_name == "tests.test_specialize"
+    assert fn_name == "cacheable_kernel"
+
+    compile_count = 0
+
+    def count_hook(*args, **kwargs):
+        nonlocal compile_count
+        compile_count += 1
+
+    triton.knobs.runtime.jit_cache_hook = count_hook
+    # clear the cache
+    specialized_kernel.device_caches.clear()
+
+    # retrieve the kernel from name and preload it.
+    fn = retrieve_fn(module_name, fn_name)
+    assert fn == specialized_kernel
+    preload = fn.preload(specialization_data)
+    assert compile_count == 1
+    assert preload.hash == hash
+
+    # verify that we hit the cache.
+    compile_count = 0
+    specialized_kernel[(1, )](o, )
+    assert compile_count == 0