From eb2cad4d8fd659a9c1872e97e2c7872c52d0b768 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 07:35:31 -0700 Subject: [PATCH 01/25] Honor single-trial autotuning for correctness tests Limit budget-one searches to one successful measurement, retry failed configurations, and skip refinement and winner runoff. Keep larger budgets unchanged and verify measurement counts through the public loader without GPU kernels. --- finegrained-kernels/README.md | 7 +- finegrained-kernels/pytest.ini | 1 + finegrained-kernels/tests/conftest.py | 5 +- .../tests/test_autotuner_budget.py | 113 ++++++++++++++++++ .../finegrained_kernels/bayesian_autotuner.py | 10 +- 5 files changed, 129 insertions(+), 7 deletions(-) create mode 100644 finegrained-kernels/tests/test_autotuner_budget.py diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index a4e1bfb1f..d1bc82e0c 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,8 +157,11 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget; `FINEGRAINED_AUTOTUNE_LOG=` appends -per-config timings as JSONL to ``. +`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. Set it to `1` for correctness tests: +the tuner stops after one successful measurement, retrying failed candidates, and skips the +additional basin probes, coordinate refinement, and winner remeasurement. Larger budgets retain +those performance search steps. `FINEGRAINED_AUTOTUNE_LOG=` appends per-config timings as +JSONL to ``. ## Tests diff --git a/finegrained-kernels/pytest.ini b/finegrained-kernels/pytest.ini index 4abbee69d..7cbb14e8a 100644 --- a/finegrained-kernels/pytest.ini +++ b/finegrained-kernels/pytest.ini @@ -9,6 +9,7 @@ testpaths = tests addopts = --dist worksteal markers = + cpu: infrastructure tests that do not launch device kernels kernels_ci: mark test to run in kernels CI environment benchmark: mark test as a performance benchmark (run with -s to see timings) slow: long-running test (multi-process determinism checks, the forced-config sweep); deselect with -m "not slow" diff --git a/finegrained-kernels/tests/conftest.py b/finegrained-kernels/tests/conftest.py index 7eaf9e726..326430f90 100644 --- a/finegrained-kernels/tests/conftest.py +++ b/finegrained-kernels/tests/conftest.py @@ -99,11 +99,12 @@ def _pin_worker_gpu() -> None: def pytest_collection_modifyitems(config, items): - """Skip every test when FP8 is unsupported — kernels require SM90+ on CUDA + """Skip device tests when FP8 is unsupported — kernels require SM90+ on CUDA (Hopper/Blackwell) or XPU. SM<89 can't even compile ``fp8e4nv``; SM89 (Ada) compiles but has numerics drift we don't gate against.""" if SUPPORTS_FP8: return skip = pytest.mark.skip(reason="FP8 kernels require SM90+ on CUDA or XPU") for item in items: - item.add_marker(skip) + if item.get_closest_marker("cpu") is None: + item.add_marker(skip) diff --git a/finegrained-kernels/tests/test_autotuner_budget.py b/finegrained-kernels/tests/test_autotuner_budget.py new file mode 100644 index 000000000..6c50a621a --- /dev/null +++ b/finegrained-kernels/tests/test_autotuner_budget.py @@ -0,0 +1,113 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Exercise real search and benchmarking plumbing without launching device kernels.""" + +import importlib +from collections import Counter + +import kernels +import pytest +import triton + + +pytestmark = [pytest.mark.cpu, pytest.mark.kernels_ci] + + +@pytest.fixture(scope="module") +def autotuner_class(): + kernel = kernels.get_kernel( + "kernels-community/finegrained-kernels", version=0, check_arch=False + ) + return importlib.import_module( + kernel.__name__ + ".bayesian_autotuner" + ).BayesianAutotuner + + +class HostKernel: + def __init__(self): + # Triton's constructor unwraps .fn until it reaches a Python function. + self.fn = lambda: None + self.__name__ = "host_budget_probe" + self.launches = [] + + def run(self, **kwargs): + self.launches.append((kwargs["MODE"], kwargs["BLOCK_SIZE_N"])) + + +def make_tuner(autotuner_class, first_failure=None): + kernel = HostKernel() + attempts, successful = [], [] + + def measure(kernel_call, quantiles): + attempts.append(None) + if len(attempts) == 1 and first_failure == "exception": + raise RuntimeError("synthetic rejected configuration") + kernel_call() + if len(attempts) == 1 and first_failure == "inf": + return [float("inf")] * 3 + mode, tile = kernel.launches[-1] + successful.append((mode, tile)) + return [1.0 / tile] * 3 + + configs = [ + triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}, num_warps=4, num_stages=2) + for mode in ("plain", "swap", "tma") + for tile in (16, 32, 64) + ] + tuner = autotuner_class( + kernel, + [], + configs, + [], + None, + None, + n_trials=80, + refine=True, + path_anchor_axes=("MODE",), + cache_results=False, + do_bench=measure, + ) + return tuner, kernel, attempts, successful + + +@pytest.mark.parametrize("first_failure", [None, "exception", "inf"]) +def test_single_trial_measures_one_success(autotuner_class, monkeypatch, first_failure): + """Multiple basins and enabled refinement must not expand the one-trial budget.""" + monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "1") + tuner, kernel, attempts, successful = make_tuner(autotuner_class, first_failure) + tuner.run() + assert len(successful) == 1 + assert len(attempts) == (1 if first_failure is None else 2) + assert kernel.launches[-1] == successful[0] + # A repeated key is a normal cached launch, not a second measurement. + tuner.run() + assert len(successful) == 1 + assert len(attempts) == (1 if first_failure is None else 2) + + +def test_larger_budget_keeps_anchors_refinement_and_runoff( + autotuner_class, monkeypatch +): + """Budget 2 keeps the established wider search: six anchors, a neighbor and runoff.""" + monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "2") + tuner, kernel, attempts, successful = make_tuner(autotuner_class) + tuner.run() + counts = Counter(successful) + assert set(mode for mode, _ in successful) == {"plain", "swap", "tma"} + assert any( + tile == 16 for _, tile in successful + ) # refinement beyond median/max anchors + assert all(counts[(mode, 64)] == 2 for mode in ("plain", "swap", "tma")) # runoff + assert len(successful) == len(attempts) == 10 + assert kernel.launches[-1][1] == 64 diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 31a0aa68f..2f5a42c09 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -489,6 +489,7 @@ def bench_idx(idx: int) -> float: # per-dimension model can write off a whole axis it never saw succeed (two dsv4 tunes # shipped 25-60% slow winners because their random seeds only sampled a basin's dead # configs), and coordinate descent can't recover a winner two coupled flips away. + single_trial = self.n_trials == 1 n_startup = max(2, min(self.n_startup_trials, self.n_trials)) anchors = self._basin_anchor_indices(configs) order = list(range(len(configs))) @@ -496,7 +497,10 @@ def bench_idx(idx: int) -> float: warm_idx = self._warm_start_index(configs, key) head = anchors + ([warm_idx] if warm_idx is not None else []) order = list(dict.fromkeys(head + order)) - for idx in order[: max(n_startup, len(head))]: + # Budget 1 is the fast correctness-test mode: stop after one valid measurement, + # rather than probing every basin. The TPE loop below retries failed candidates. + seed_count = 1 if single_trial else max(n_startup, len(head)) + for idx in order[:seed_count]: bench_idx(idx) # TPE: split measured configs into good (top-gamma) / bad, build per-dimension value @@ -574,7 +578,7 @@ def bench_idx(idx: int) -> float: # Coordinate-descent refinement: try single-dim perturbations around # the current best until no neighbor improves. - if self.refine: + if self.refine and not single_trial: for _ in range(self.max_refine_iters): best_idx = min(timings, key=timings.get) best_sig = sigs[best_idx] @@ -597,7 +601,7 @@ def bench_idx(idx: int) -> float: # the top candidates back-to-back — same warmed process, adjacent in time, so they see # the same clocks — and crown the winner of THAT comparison. Costs ~3 extra benches. finite = sorted((t, i) for i, t in timings.items() if t != float("inf")) - if len(finite) > 1: + if len(finite) > 1 and not single_trial: for _, i in finite[:3]: timings.pop(i, None) bench_idx(i) # memoizing closure: re-benches and re-records the popped entry From 58e6a51a6a7cae840384803e438d873f46a1d6df Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 11:43:21 -0700 Subject: [PATCH 02/25] Keep trial-budget regressions in the existing direct-import suite --- finegrained-kernels/pytest.ini | 1 - finegrained-kernels/tests/conftest.py | 5 +- finegrained-kernels/tests/test_autotuner.py | 82 ++++++++++++- .../tests/test_autotuner_budget.py | 113 ------------------ 4 files changed, 83 insertions(+), 118 deletions(-) delete mode 100644 finegrained-kernels/tests/test_autotuner_budget.py diff --git a/finegrained-kernels/pytest.ini b/finegrained-kernels/pytest.ini index 7cbb14e8a..4abbee69d 100644 --- a/finegrained-kernels/pytest.ini +++ b/finegrained-kernels/pytest.ini @@ -9,7 +9,6 @@ testpaths = tests addopts = --dist worksteal markers = - cpu: infrastructure tests that do not launch device kernels kernels_ci: mark test to run in kernels CI environment benchmark: mark test as a performance benchmark (run with -s to see timings) slow: long-running test (multi-process determinism checks, the forced-config sweep); deselect with -m "not slow" diff --git a/finegrained-kernels/tests/conftest.py b/finegrained-kernels/tests/conftest.py index 326430f90..7eaf9e726 100644 --- a/finegrained-kernels/tests/conftest.py +++ b/finegrained-kernels/tests/conftest.py @@ -99,12 +99,11 @@ def _pin_worker_gpu() -> None: def pytest_collection_modifyitems(config, items): - """Skip device tests when FP8 is unsupported — kernels require SM90+ on CUDA + """Skip every test when FP8 is unsupported — kernels require SM90+ on CUDA (Hopper/Blackwell) or XPU. SM<89 can't even compile ``fp8e4nv``; SM89 (Ada) compiles but has numerics drift we don't gate against.""" if SUPPORTS_FP8: return skip = pytest.mark.skip(reason="FP8 kernels require SM90+ on CUDA or XPU") for item in items: - if item.get_closest_marker("cpu") is None: - item.add_marker(skip) + item.add_marker(skip) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 38ddc7326..331d5eeed 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -17,6 +17,7 @@ carries its incident.""" import logging +from collections import Counter import subprocess import sys @@ -29,7 +30,7 @@ import finegrained_kernels # type: ignore import finegrained_kernels.matmul # type: ignore -from finegrained_kernels.bayesian_autotuner import bayesian_autotune # type: ignore +from finegrained_kernels.bayesian_autotuner import BayesianAutotuner, bayesian_autotune # type: ignore from finegrained_kernels.compat import is_sm10x # type: ignore from finegrained_kernels.pruners import block_within_dim_pruner, mx_config_pruner # type: ignore @@ -343,3 +344,82 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): assert len(msgs) == 2, msgs assert "[memoized]" not in msgs[0] or "[memoized]" in msgs[1] assert "[memoized]" in msgs[1], msgs[1] + + +class HostKernel: + def __init__(self): + # Triton's constructor unwraps .fn until it reaches a Python function. + self.fn = lambda: None + self.__name__ = "host_budget_probe" + self.launches = [] + + def run(self, **kwargs): + self.launches.append((kwargs["MODE"], kwargs["BLOCK_SIZE_N"])) + + +def make_tuner(first_failure=None): + kernel = HostKernel() + attempts, successful = [], [] + + def measure(kernel_call, quantiles): + attempts.append(None) + if len(attempts) == 1 and first_failure == "exception": + raise RuntimeError("synthetic rejected configuration") + kernel_call() + if len(attempts) == 1 and first_failure == "inf": + return [float("inf")] * 3 + mode, tile = kernel.launches[-1] + successful.append((mode, tile)) + return [1.0 / tile] * 3 + + configs = [ + triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}, num_warps=4, num_stages=2) + for mode in ("plain", "swap", "tma") + for tile in (16, 32, 64) + ] + tuner = BayesianAutotuner( + kernel, + [], + configs, + [], + None, + None, + n_trials=80, + refine=True, + path_anchor_axes=("MODE",), + cache_results=False, + do_bench=measure, + ) + return tuner, kernel, attempts, successful + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("first_failure", [None, "exception", "inf"]) +def test_single_trial_measures_one_success(monkeypatch, first_failure): + """Multiple basins and enabled refinement must not expand the one-trial budget.""" + monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "1") + tuner, kernel, attempts, successful = make_tuner(first_failure) + tuner.run() + assert len(successful) == 1 + assert len(attempts) == (1 if first_failure is None else 2) + assert kernel.launches[-1] == successful[0] + # A repeated key is a normal cached launch, not a second measurement. + tuner.run() + assert len(successful) == 1 + assert len(attempts) == (1 if first_failure is None else 2) + + +@pytest.mark.kernels_ci +def test_larger_budget_keeps_anchors_refinement_and_runoff(monkeypatch): + """Budget 2 keeps the established wider search: six anchors, a neighbor and runoff.""" + monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "2") + tuner, kernel, attempts, successful = make_tuner() + tuner.run() + counts = Counter(successful) + assert set(mode for mode, _ in successful) == {"plain", "swap", "tma"} + assert any( + tile == 16 for _, tile in successful + ) # refinement beyond median/max anchors + assert all(counts[(mode, 64)] == 2 for mode in ("plain", "swap", "tma")) # runoff + assert len(successful) == len(attempts) == 10 + assert kernel.launches[-1][1] == 64 diff --git a/finegrained-kernels/tests/test_autotuner_budget.py b/finegrained-kernels/tests/test_autotuner_budget.py deleted file mode 100644 index 6c50a621a..000000000 --- a/finegrained-kernels/tests/test_autotuner_budget.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Exercise real search and benchmarking plumbing without launching device kernels.""" - -import importlib -from collections import Counter - -import kernels -import pytest -import triton - - -pytestmark = [pytest.mark.cpu, pytest.mark.kernels_ci] - - -@pytest.fixture(scope="module") -def autotuner_class(): - kernel = kernels.get_kernel( - "kernels-community/finegrained-kernels", version=0, check_arch=False - ) - return importlib.import_module( - kernel.__name__ + ".bayesian_autotuner" - ).BayesianAutotuner - - -class HostKernel: - def __init__(self): - # Triton's constructor unwraps .fn until it reaches a Python function. - self.fn = lambda: None - self.__name__ = "host_budget_probe" - self.launches = [] - - def run(self, **kwargs): - self.launches.append((kwargs["MODE"], kwargs["BLOCK_SIZE_N"])) - - -def make_tuner(autotuner_class, first_failure=None): - kernel = HostKernel() - attempts, successful = [], [] - - def measure(kernel_call, quantiles): - attempts.append(None) - if len(attempts) == 1 and first_failure == "exception": - raise RuntimeError("synthetic rejected configuration") - kernel_call() - if len(attempts) == 1 and first_failure == "inf": - return [float("inf")] * 3 - mode, tile = kernel.launches[-1] - successful.append((mode, tile)) - return [1.0 / tile] * 3 - - configs = [ - triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}, num_warps=4, num_stages=2) - for mode in ("plain", "swap", "tma") - for tile in (16, 32, 64) - ] - tuner = autotuner_class( - kernel, - [], - configs, - [], - None, - None, - n_trials=80, - refine=True, - path_anchor_axes=("MODE",), - cache_results=False, - do_bench=measure, - ) - return tuner, kernel, attempts, successful - - -@pytest.mark.parametrize("first_failure", [None, "exception", "inf"]) -def test_single_trial_measures_one_success(autotuner_class, monkeypatch, first_failure): - """Multiple basins and enabled refinement must not expand the one-trial budget.""" - monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "1") - tuner, kernel, attempts, successful = make_tuner(autotuner_class, first_failure) - tuner.run() - assert len(successful) == 1 - assert len(attempts) == (1 if first_failure is None else 2) - assert kernel.launches[-1] == successful[0] - # A repeated key is a normal cached launch, not a second measurement. - tuner.run() - assert len(successful) == 1 - assert len(attempts) == (1 if first_failure is None else 2) - - -def test_larger_budget_keeps_anchors_refinement_and_runoff( - autotuner_class, monkeypatch -): - """Budget 2 keeps the established wider search: six anchors, a neighbor and runoff.""" - monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "2") - tuner, kernel, attempts, successful = make_tuner(autotuner_class) - tuner.run() - counts = Counter(successful) - assert set(mode for mode, _ in successful) == {"plain", "swap", "tma"} - assert any( - tile == 16 for _, tile in successful - ) # refinement beyond median/max anchors - assert all(counts[(mode, 64)] == 2 for mode in ("plain", "swap", "tma")) # runoff - assert len(successful) == len(attempts) == 10 - assert kernel.launches[-1][1] == 64 From cc3656fc412c51d28598f9bf502c95c37d315e61 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 12:02:08 -0700 Subject: [PATCH 03/25] Reduce single-trial coverage to one focused benchmark-count regression --- finegrained-kernels/tests/test_autotuner.py | 93 ++++----------------- 1 file changed, 17 insertions(+), 76 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 331d5eeed..49eab6098 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -17,9 +17,9 @@ carries its incident.""" import logging -from collections import Counter import subprocess import sys +from unittest.mock import Mock import pytest import torch @@ -30,7 +30,7 @@ import finegrained_kernels # type: ignore import finegrained_kernels.matmul # type: ignore -from finegrained_kernels.bayesian_autotuner import BayesianAutotuner, bayesian_autotune # type: ignore +from finegrained_kernels.bayesian_autotuner import bayesian_autotune # type: ignore from finegrained_kernels.compat import is_sm10x # type: ignore from finegrained_kernels.pruners import block_within_dim_pruner, mx_config_pruner # type: ignore @@ -346,80 +346,21 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): assert "[memoized]" in msgs[1], msgs[1] -class HostKernel: - def __init__(self): - # Triton's constructor unwraps .fn until it reaches a Python function. - self.fn = lambda: None - self.__name__ = "host_budget_probe" - self.launches = [] - - def run(self, **kwargs): - self.launches.append((kwargs["MODE"], kwargs["BLOCK_SIZE_N"])) - - -def make_tuner(first_failure=None): - kernel = HostKernel() - attempts, successful = [], [] - - def measure(kernel_call, quantiles): - attempts.append(None) - if len(attempts) == 1 and first_failure == "exception": - raise RuntimeError("synthetic rejected configuration") - kernel_call() - if len(attempts) == 1 and first_failure == "inf": - return [float("inf")] * 3 - mode, tile = kernel.launches[-1] - successful.append((mode, tile)) - return [1.0 / tile] * 3 - - configs = [ - triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}, num_warps=4, num_stages=2) - for mode in ("plain", "swap", "tma") - for tile in (16, 32, 64) - ] - tuner = BayesianAutotuner( - kernel, - [], - configs, - [], - None, - None, - n_trials=80, - refine=True, - path_anchor_axes=("MODE",), - cache_results=False, - do_bench=measure, - ) - return tuner, kernel, attempts, successful - - @pytest.mark.kernels_ci -@pytest.mark.parametrize("first_failure", [None, "exception", "inf"]) -def test_single_trial_measures_one_success(monkeypatch, first_failure): - """Multiple basins and enabled refinement must not expand the one-trial budget.""" +def test_single_trial_measures_one_config(monkeypatch): monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "1") - tuner, kernel, attempts, successful = make_tuner(first_failure) - tuner.run() - assert len(successful) == 1 - assert len(attempts) == (1 if first_failure is None else 2) - assert kernel.launches[-1] == successful[0] - # A repeated key is a normal cached launch, not a second measurement. - tuner.run() - assert len(successful) == 1 - assert len(attempts) == (1 if first_failure is None else 2) - + configs = [ + triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}) + for mode in ("plain", "swap") for tile in (16, 32, 64) + ] -@pytest.mark.kernels_ci -def test_larger_budget_keeps_anchors_refinement_and_runoff(monkeypatch): - """Budget 2 keeps the established wider search: six anchors, a neighbor and runoff.""" - monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "2") - tuner, kernel, attempts, successful = make_tuner() - tuner.run() - counts = Counter(successful) - assert set(mode for mode, _ in successful) == {"plain", "swap", "tma"} - assert any( - tile == 16 for _, tile in successful - ) # refinement beyond median/max anchors - assert all(counts[(mode, 64)] == 2 for mode in ("plain", "swap", "tma")) # runoff - assert len(successful) == len(attempts) == 10 - assert kernel.launches[-1][1] == 64 + @bayesian_autotune(configs, [], n_trials=80, refine=True, path_anchor_axes=("MODE",)) + @triton.jit + def kernel(MODE: tl.constexpr, BLOCK_SIZE_N: tl.constexpr): + pass + + measure = Mock(return_value=1.0) + monkeypatch.setattr(kernel, "_bench", measure) + winner = kernel._bayesian_search(configs, (), {}, ()) + measure.assert_called_once() + assert winner is measure.call_args.kwargs["config"] From 2b928b33fe4dcb22fbfe15c9722028d8b51033bd Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 12:07:19 -0700 Subject: [PATCH 04/25] Control autotuning search phases explicitly --- finegrained-kernels/README.md | 11 ++++--- finegrained-kernels/tests/test_autotuner.py | 9 ++++-- .../finegrained_kernels/bayesian_autotuner.py | 31 ++++++++++++------- 3 files changed, 32 insertions(+), 19 deletions(-) diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index d1bc82e0c..f02343d52 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,11 +157,12 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. Set it to `1` for correctness tests: -the tuner stops after one successful measurement, retrying failed candidates, and skips the -additional basin probes, coordinate refinement, and winner remeasurement. Larger budgets retain -those performance search steps. `FINEGRAINED_AUTOTUNE_LOG=` appends per-config timings as -JSONL to ``. +`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. Basin probes, coordinate refinement, +and winner remeasurement are independently controlled by `FINEGRAINED_AUTOTUNE_BASIN_PROBES`, +`FINEGRAINED_AUTOTUNE_REFINE`, and `FINEGRAINED_AUTOTUNE_RUNOFF` (`0` or `1`, all enabled by default). +For correctness tests, set the trial budget to `1` and all three controls to `0`: only one +successful configuration is measured, retrying failed candidates. `FINEGRAINED_AUTOTUNE_LOG=` +appends per-config timings as JSONL to ``. ## Tests diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 49eab6098..430d76291 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -347,8 +347,11 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): @pytest.mark.kernels_ci -def test_single_trial_measures_one_config(monkeypatch): +@pytest.mark.parametrize("search_phases", [False, True]) +def test_autotune_search_controls(monkeypatch, search_phases): monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "1") + for phase in ("BASIN_PROBES", "REFINE", "RUNOFF"): + monkeypatch.setenv(f"FINEGRAINED_AUTOTUNE_{phase}", str(int(search_phases))) configs = [ triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}) for mode in ("plain", "swap") for tile in (16, 32, 64) @@ -362,5 +365,5 @@ def kernel(MODE: tl.constexpr, BLOCK_SIZE_N: tl.constexpr): measure = Mock(return_value=1.0) monkeypatch.setattr(kernel, "_bench", measure) winner = kernel._bayesian_search(configs, (), {}, ()) - measure.assert_called_once() - assert winner is measure.call_args.kwargs["config"] + assert (measure.call_count > 1) == search_phases + assert winner in [call.kwargs["config"] for call in measure.call_args_list] diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 2f5a42c09..1d99f2535 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -84,7 +84,9 @@ def __init__( max_failures: int | None = None, n_startup_trials: int = 12, gamma: float = 0.25, + basin_probes: bool = True, refine: bool = True, + runoff: bool = True, max_refine_iters: int = 5, log_path: str | None = None, path_anchor_axes: tuple[str, ...] = (), @@ -110,7 +112,9 @@ def __init__( self.n_startup_trials = n_startup_trials # top fraction of measured configs the TPE treats as "good" self.gamma = gamma - self.refine = refine + self.basin_probes = os.environ.get("FINEGRAINED_AUTOTUNE_BASIN_PROBES", str(int(basin_probes))) == "1" + self.refine = os.environ.get("FINEGRAINED_AUTOTUNE_REFINE", str(int(refine))) == "1" + self.runoff = os.environ.get("FINEGRAINED_AUTOTUNE_RUNOFF", str(int(runoff))) == "1" self.max_refine_iters = max_refine_iters # Branch axes (by kwarg name) that relocate the TILE optimum — the caller's # declaration, since it is configuration knowledge (the tuner never hardcodes @@ -489,18 +493,14 @@ def bench_idx(idx: int) -> float: # per-dimension model can write off a whole axis it never saw succeed (two dsv4 tunes # shipped 25-60% slow winners because their random seeds only sampled a basin's dead # configs), and coordinate descent can't recover a winner two coupled flips away. - single_trial = self.n_trials == 1 - n_startup = max(2, min(self.n_startup_trials, self.n_trials)) - anchors = self._basin_anchor_indices(configs) + n_startup = min(self.n_trials, max(2, self.n_startup_trials)) + anchors = self._basin_anchor_indices(configs) if self.basin_probes else [] order = list(range(len(configs))) random.Random(0).shuffle(order) warm_idx = self._warm_start_index(configs, key) head = anchors + ([warm_idx] if warm_idx is not None else []) order = list(dict.fromkeys(head + order)) - # Budget 1 is the fast correctness-test mode: stop after one valid measurement, - # rather than probing every basin. The TPE loop below retries failed candidates. - seed_count = 1 if single_trial else max(n_startup, len(head)) - for idx in order[:seed_count]: + for idx in order[: max(n_startup, len(head))]: bench_idx(idx) # TPE: split measured configs into good (top-gamma) / bad, build per-dimension value @@ -578,7 +578,7 @@ def bench_idx(idx: int) -> float: # Coordinate-descent refinement: try single-dim perturbations around # the current best until no neighbor improves. - if self.refine and not single_trial: + if self.refine: for _ in range(self.max_refine_iters): best_idx = min(timings, key=timings.get) best_sig = sigs[best_idx] @@ -601,7 +601,7 @@ def bench_idx(idx: int) -> float: # the top candidates back-to-back — same warmed process, adjacent in time, so they see # the same clocks — and crown the winner of THAT comparison. Costs ~3 extra benches. finite = sorted((t, i) for i, t in timings.items() if t != float("inf")) - if len(finite) > 1 and not single_trial: + if self.runoff and len(finite) > 1: for _, i in finite[:3]: timings.pop(i, None) bench_idx(i) # memoizing closure: re-benches and re-records the popped entry @@ -760,7 +760,8 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): # and FINEGRAINED_AUTOTUNE_TRIALS is not in Triton's invalidating-env list — # without n_trials a 3-trial exploratory winner gets replayed at 100 trials. str(self.keys), - f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}", + f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," + f"basin_probes={self.basin_probes},refine={self.refine},runoff={self.runoff}", str(tuning_key), ] + [str(c) for c in configs] # str(Config) is pre_hook-free and process-stable cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest() @@ -896,7 +897,9 @@ def bayesian_autotune( max_failures: int | None = None, n_startup_trials: int = 12, gamma: float = 0.25, + basin_probes: bool = True, refine: bool = True, + runoff: bool = True, max_refine_iters: int = 5, log_path: str | None = None, path_anchor_axes: tuple[str, ...] = (), @@ -911,7 +914,11 @@ def bayesian_autotune( that fail to compile/run are skipped without consuming it) n_startup_trials: random seed configs before the TPE model kicks in gamma: top fraction of measured configs treated as "good" + basin_probes: measure categorical basin anchors before the TPE refine, max_refine_iters: coordinate-descent refinement after the TPE + runoff: remeasure the top candidates before choosing the winner + (basin_probes/refine/runoff each have an env override: + FINEGRAINED_AUTOTUNE_BASIN_PROBES/_REFINE/_RUNOFF, 0 or 1) path_anchor_axes: branch axes (kwarg names) whose values relocate the tile optimum — one guaranteed max-tile anchor per value combo (see _basin_anchor_indices); the kernel declares them, the @@ -937,7 +944,9 @@ def decorator(fn): max_failures=max_failures, n_startup_trials=n_startup_trials, gamma=gamma, + basin_probes=basin_probes, refine=refine, + runoff=runoff, max_refine_iters=max_refine_iters, log_path=log_path, path_anchor_axes=path_anchor_axes, From 8f876b05fdea276db61a75eb55bf55cc92283fbc Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 15:48:57 -0700 Subject: [PATCH 05/25] Check autotuning phase controls independently --- finegrained-kernels/tests/test_autotuner.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 430d76291..455f90541 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -347,11 +347,13 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): @pytest.mark.kernels_ci -@pytest.mark.parametrize("search_phases", [False, True]) -def test_autotune_search_controls(monkeypatch, search_phases): - monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", "1") - for phase in ("BASIN_PROBES", "REFINE", "RUNOFF"): - monkeypatch.setenv(f"FINEGRAINED_AUTOTUNE_{phase}", str(int(search_phases))) +@pytest.mark.parametrize("phase", ["BASIN_PROBES", "REFINE", "RUNOFF"]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_autotune_search_controls(monkeypatch, phase, enabled): + budget = 2 if phase == "RUNOFF" else 1 + monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", str(budget)) + for name in ("BASIN_PROBES", "REFINE", "RUNOFF"): + monkeypatch.setenv(f"FINEGRAINED_AUTOTUNE_{name}", str(int(enabled and name == phase))) configs = [ triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}) for mode in ("plain", "swap") for tile in (16, 32, 64) @@ -365,5 +367,5 @@ def kernel(MODE: tl.constexpr, BLOCK_SIZE_N: tl.constexpr): measure = Mock(return_value=1.0) monkeypatch.setattr(kernel, "_bench", measure) winner = kernel._bayesian_search(configs, (), {}, ()) - assert (measure.call_count > 1) == search_phases + assert (measure.call_count > budget) == enabled assert winner in [call.kwargs["config"] for call in measure.call_args_list] From e3a6cceac1fb734a8cd45a35dcda4e476caede37 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:51:22 -0700 Subject: [PATCH 06/25] Drop tests for autotuning environment controls --- finegrained-kernels/tests/test_autotuner.py | 26 --------------------- 1 file changed, 26 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 455f90541..38ddc7326 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -19,7 +19,6 @@ import logging import subprocess import sys -from unittest.mock import Mock import pytest import torch @@ -344,28 +343,3 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): assert len(msgs) == 2, msgs assert "[memoized]" not in msgs[0] or "[memoized]" in msgs[1] assert "[memoized]" in msgs[1], msgs[1] - - -@pytest.mark.kernels_ci -@pytest.mark.parametrize("phase", ["BASIN_PROBES", "REFINE", "RUNOFF"]) -@pytest.mark.parametrize("enabled", [False, True]) -def test_autotune_search_controls(monkeypatch, phase, enabled): - budget = 2 if phase == "RUNOFF" else 1 - monkeypatch.setenv("FINEGRAINED_AUTOTUNE_TRIALS", str(budget)) - for name in ("BASIN_PROBES", "REFINE", "RUNOFF"): - monkeypatch.setenv(f"FINEGRAINED_AUTOTUNE_{name}", str(int(enabled and name == phase))) - configs = [ - triton.Config({"MODE": mode, "BLOCK_SIZE_N": tile}) - for mode in ("plain", "swap") for tile in (16, 32, 64) - ] - - @bayesian_autotune(configs, [], n_trials=80, refine=True, path_anchor_axes=("MODE",)) - @triton.jit - def kernel(MODE: tl.constexpr, BLOCK_SIZE_N: tl.constexpr): - pass - - measure = Mock(return_value=1.0) - monkeypatch.setattr(kernel, "_bench", measure) - winner = kernel._bayesian_search(configs, (), {}, ()) - assert (measure.call_count > budget) == enabled - assert winner in [call.kwargs["config"] for call in measure.call_args_list] From 23265027acf09a63c7e9b51ee01f90e774a21fa0 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 15:48:57 -0700 Subject: [PATCH 07/25] Select interior autotuning anchors across ordinal axes --- finegrained-kernels/tests/test_autotuner.py | 25 ++++++- .../finegrained_kernels/bayesian_autotuner.py | 67 +++++++++---------- 2 files changed, 57 insertions(+), 35 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 38ddc7326..c2e3285b6 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -29,7 +29,7 @@ import finegrained_kernels # type: ignore import finegrained_kernels.matmul # type: ignore -from finegrained_kernels.bayesian_autotuner import bayesian_autotune # type: ignore +from finegrained_kernels.bayesian_autotuner import BayesianAutotuner, bayesian_autotune # type: ignore from finegrained_kernels.compat import is_sm10x # type: ignore from finegrained_kernels.pruners import block_within_dim_pruner, mx_config_pruner # type: ignore @@ -343,3 +343,26 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): assert len(msgs) == 2, msgs assert "[memoized]" not in msgs[0] or "[memoized]" in msgs[1] assert "[memoized]" in msgs[1], msgs[1] + + +@pytest.mark.kernels_ci +def test_basin_anchors_cover_interior_and_wide_output_tiles(): + tuner = object.__new__(BayesianAutotuner) + tuner.path_anchor_axes = ("SWAP_AB",) + configs = [ + triton.Config( + {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "SWAP_AB": swap}, + num_warps=warps, num_stages=stages, + ) + for swap in (False, True) + for bm in (16, 64, 128) for bn in (64, 128, 256) for bk in (64, 128, 256) + for warps in (4, 8, 16) for stages in (2, 4, 6) + ] + anchors = [configs[i] for i in tuner._basin_anchor_indices(configs)] + assert len(anchors) == 4 + for c in anchors: + assert c.kwargs["BLOCK_SIZE_K"] == 128 + assert c.num_warps == 8 and c.num_stages == 4 + assert [(c.kwargs["BLOCK_SIZE_M"], c.kwargs["BLOCK_SIZE_N"]) for c in anchors] == [ + (64, 128), (128, 256), (64, 128), (128, 256), + ] diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 1d99f2535..4e439fb87 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -55,6 +55,9 @@ logger = logging.getLogger(__name__) +# Retune crowns produced before the ordinal/interior anchor policy. +_TUNING_POLICY = "ordinal-interior-anchors" + # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every # BayesianAutotuner.run() call increments its kernel's tally, dumped to at exit. # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token @@ -643,12 +646,8 @@ def _log_result(self, key, config: Config, ms: float): pass def _basin_anchor_indices(self, configs: List[Config]) -> List[int]: - """One representative config index per (COMPUTE_MODE, SWAP_AB) basin — the MEDIAN in - tile-sort order, a mid-sized tile with mid warps/stages. Not the smallest: a basin's - minimal corner (min BK x 2 warps) can be latency-bound pathological (a 131µs anchor in - a basin whose peak is 41µs re-poisons the axis it was meant to protect). Coordinate - descent climbs BN/BK/warps/stages from wherever the TPE lands within the basin. - Returns [] when the grid has no such axes (single basin).""" + """Interior representatives of categorical basins, plus a wide output tile + per declared compute path. All anchors come from the pruned grid.""" # Basin axes are DERIVED, not declared: a config kwarg with string or boolean # values is a branch axis (different code path — compute mode, operand swap, # warp specialization), which partitions the grid into disjoint performance @@ -674,41 +673,40 @@ def _basin_anchor_indices(self, configs: List[Config]) -> List[int]: if len(groups) <= 1: return [] - def tile_order(i): - return ( - configs[i].kwargs.get("BLOCK_SIZE_N", 0), - configs[i].kwargs.get("BLOCK_SIZE_K", 0), - configs[i].num_warps, - configs[i].num_stages, - ) - - # Median anchor per basin, PLUS one max-tile anchor per declared PATH-axis combo - # (``path_anchor_axes``). The median alone left the big-tile end unsampled: TPE - # densities lock onto the first-timed mid-tile region and coordinate descent cannot - # cross DIAGONAL ridges — dsv4 mxfp8 prefill gate_up's BN=128/BK=256 winner (840us) - # needs BN and warps to move together, so descent from the BN=64 crown rejects the - # lone BN step and n_trials 100/200/400 all missed it; one guaranteed big-tile - # sample fixed it (crown 958 -> 880 at the default budget). The max anchor rides - # the declared COARSE grouping, NOT the full branch-axis cross: per-basin max - # anchors on the weight-only grouped grid (32 basins: modes x WS x PS x 2 memory - # modes each way) put 64 forced trials into a 100-trial budget and starved the - # search — gpt-oss prefill regressed 3463 -> 4099us on a quiet box before this was - # scoped. Which axes qualify is the KERNEL's declaration (configuration knowledge); - # axes absent from this grid are ignored. anchors = [] coarse_seen = set() coarse_axes = [k for k in self.path_anchor_axes if k in basin_axes] for basin, idxs in groups.items(): - ordered = sorted(idxs, key=tile_order) - anchors.append(ordered[len(ordered) // 2]) + # Choose the middle of every ordinal axis together. Lexicographic sorting + # ignores M and can put the "median" at extreme K/warps/stages. + values = [configs[i].all_kwargs() for i in idxs] + axes = { + k: sorted({v[k] for v in values}) + for k, value in values[0].items() + if isinstance(value, (int, float)) and not isinstance(value, bool) + } + + def distance(i): + args = configs[i].all_kwargs() + return sum( + (options.index(args[k]) / max(1, len(options) - 1) - 0.5) ** 2 + for k, options in axes.items() + ) + + anchors.append(min(idxs, key=distance)) if not self.path_anchor_axes: - continue # no declaration -> median anchors only - # declared axes absent from THIS grid still coarse-group to () — one global - # max-tile anchor, the big-tile guarantee without per-basin anchor spam + continue coarse = tuple(basin[basin_axes.index(k)] for k in coarse_axes) - if coarse not in coarse_seen and ordered[-1] != anchors[-1]: + if coarse not in coarse_seen: coarse_seen.add(coarse) - anchors.append(ordered[-1]) + # The existing wide-tile anchor protects coupled output-tile ridges, + # with interior contraction/warps/stages rather than their largest corner. + wide = min(idxs, key=lambda i: ( + -configs[i].kwargs.get("BLOCK_SIZE_M", 1) + * configs[i].kwargs.get("BLOCK_SIZE_N", 1), distance(i), + )) + if wide != anchors[-1]: + anchors.append(wide) return anchors def _warm_start_index(self, configs: List[Config], key=None): @@ -751,6 +749,7 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): fn = fn.fn env_vars = get_cache_invalidating_env_vars() cache_key = [ + _TUNING_POLICY, triton_key(), make_backend(driver.active.get_current_target()).hash(), fn.cache_key, From c89a3e8cf77d0714711138d7f5c46eed490fc2eb Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:54:15 -0700 Subject: [PATCH 08/25] Construct anchor test tuner through its decorator --- finegrained-kernels/tests/test_autotuner.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index c2e3285b6..afc8896c7 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -29,7 +29,7 @@ import finegrained_kernels # type: ignore import finegrained_kernels.matmul # type: ignore -from finegrained_kernels.bayesian_autotuner import BayesianAutotuner, bayesian_autotune # type: ignore +from finegrained_kernels.bayesian_autotuner import bayesian_autotune # type: ignore from finegrained_kernels.compat import is_sm10x # type: ignore from finegrained_kernels.pruners import block_within_dim_pruner, mx_config_pruner # type: ignore @@ -347,8 +347,6 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): @pytest.mark.kernels_ci def test_basin_anchors_cover_interior_and_wide_output_tiles(): - tuner = object.__new__(BayesianAutotuner) - tuner.path_anchor_axes = ("SWAP_AB",) configs = [ triton.Config( {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "SWAP_AB": swap}, @@ -358,7 +356,15 @@ def test_basin_anchors_cover_interior_and_wide_output_tiles(): for bm in (16, 64, 128) for bn in (64, 128, 256) for bk in (64, 128, 256) for warps in (4, 8, 16) for stages in (2, 4, 6) ] - anchors = [configs[i] for i in tuner._basin_anchor_indices(configs)] + @bayesian_autotune(configs, [], path_anchor_axes=("SWAP_AB",)) + @triton.jit + def kernel( + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, SWAP_AB: tl.constexpr, + ): + pass + + anchors = [configs[i] for i in kernel._basin_anchor_indices(configs)] assert len(anchors) == 4 for c in anchors: assert c.kwargs["BLOCK_SIZE_K"] == 128 From 84f69c69954ec46b18b966caa2d93452cbe032af Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:54:29 -0700 Subject: [PATCH 09/25] Explain basin anchor regression expectations --- finegrained-kernels/tests/test_autotuner.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index afc8896c7..377ee5284 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -347,6 +347,9 @@ def _memo_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): @pytest.mark.kernels_ci def test_basin_anchors_cover_interior_and_wide_output_tiles(): + """Each swap basin gets an interior anchor and a wide output tile, both with + middle K, warp, and stage settings rather than lexicographic extremes. + """ configs = [ triton.Config( {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "SWAP_AB": swap}, From a7f8eff265534de75d713d26b208941472efe3a0 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:55:35 -0700 Subject: [PATCH 10/25] Remove tuning policy cache tag --- .../torch-ext/finegrained_kernels/bayesian_autotuner.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 4e439fb87..b1c205606 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -55,9 +55,6 @@ logger = logging.getLogger(__name__) -# Retune crowns produced before the ordinal/interior anchor policy. -_TUNING_POLICY = "ordinal-interior-anchors" - # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every # BayesianAutotuner.run() call increments its kernel's tally, dumped to at exit. # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token @@ -749,7 +746,6 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): fn = fn.fn env_vars = get_cache_invalidating_env_vars() cache_key = [ - _TUNING_POLICY, triton_key(), make_backend(driver.active.get_current_target()).hash(), fn.cache_key, From 4efb1d4623dc328408bf06b733693c2c05bb7ef0 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:56:58 -0700 Subject: [PATCH 11/25] Include tuner source hash in tuning cache identity --- .../torch-ext/finegrained_kernels/bayesian_autotuner.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index b1c205606..a22bfbe79 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -37,6 +37,7 @@ import time from collections import defaultdict from contextlib import contextmanager +from pathlib import Path import torch @@ -55,6 +56,9 @@ logger = logging.getLogger(__name__) +# Include the tuner implementation alongside Triton's kernel-source hash. +_TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every # BayesianAutotuner.run() call increments its kernel's tally, dumped to at exit. # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token @@ -728,8 +732,8 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): """Persist the tuned best config to Triton's on-disk cache so a later run (or process) skips the whole minutes-long search+compile — most of which is Triton codegen, not benching. Keyed exactly like Triton's own autotune cache (triton version + backend + - kernel source hash + invalidating env + key + the full config grid), so any of those - changing re-tunes. + kernel/tuner source hashes + invalidating env + key + the full config grid), so any + of those changing re-tunes. Unlike Triton's stock version we do NOT bail when configs carry a ``pre_hook`` (ours always do): we persist each config's ``all_kwargs()`` — serializable and pre_hook-free — @@ -749,6 +753,7 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): triton_key(), make_backend(driver.active.get_current_target()).hash(), fn.cache_key, + _TUNER_SOURCE_HASH, str(sorted(env_vars.items())), # Key NAMES and the search budget are determinants too: tuning_key holds # values only (renaming a key or swapping two same-typed keys would collide), From be7b8965843c5ad22d4da2c550675229a8267a0e Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:59:42 -0700 Subject: [PATCH 12/25] Control refinement with an iteration limit environment variable --- finegrained-kernels/README.md | 7 +-- .../finegrained_kernels/bayesian_autotuner.py | 46 +++++++++---------- 2 files changed, 25 insertions(+), 28 deletions(-) diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index f02343d52..7398d8e9b 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,9 +157,10 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. Basin probes, coordinate refinement, -and winner remeasurement are independently controlled by `FINEGRAINED_AUTOTUNE_BASIN_PROBES`, -`FINEGRAINED_AUTOTUNE_REFINE`, and `FINEGRAINED_AUTOTUNE_RUNOFF` (`0` or `1`, all enabled by default). +`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. Basin probes and winner remeasurement +are controlled by `FINEGRAINED_AUTOTUNE_BASIN_PROBES` and `FINEGRAINED_AUTOTUNE_RUNOFF` +(`0` or `1`, both enabled by default). `FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS` controls the +coordinate-refinement iteration limit (default `5`; `0` disables refinement). For correctness tests, set the trial budget to `1` and all three controls to `0`: only one successful configuration is measured, retrying failed candidates. `FINEGRAINED_AUTOTUNE_LOG=` appends per-config timings as JSONL to ``. diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index a22bfbe79..8ec0a0e9d 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -89,7 +89,6 @@ def __init__( n_startup_trials: int = 12, gamma: float = 0.25, basin_probes: bool = True, - refine: bool = True, runoff: bool = True, max_refine_iters: int = 5, log_path: str | None = None, @@ -117,9 +116,8 @@ def __init__( # top fraction of measured configs the TPE treats as "good" self.gamma = gamma self.basin_probes = os.environ.get("FINEGRAINED_AUTOTUNE_BASIN_PROBES", str(int(basin_probes))) == "1" - self.refine = os.environ.get("FINEGRAINED_AUTOTUNE_REFINE", str(int(refine))) == "1" self.runoff = os.environ.get("FINEGRAINED_AUTOTUNE_RUNOFF", str(int(runoff))) == "1" - self.max_refine_iters = max_refine_iters + self.max_refine_iters = int(os.environ.get("FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS", str(max_refine_iters))) # Branch axes (by kwarg name) that relocate the TILE optimum — the caller's # declaration, since it is configuration knowledge (the tuner never hardcodes # axis names). Each distinct value combination of these axes gets one guaranteed @@ -582,22 +580,21 @@ def bench_idx(idx: int) -> float: # Coordinate-descent refinement: try single-dim perturbations around # the current best until no neighbor improves. - if self.refine: - for _ in range(self.max_refine_iters): - best_idx = min(timings, key=timings.get) - best_sig = sigs[best_idx] - best_ms = timings[best_idx] - improved = False - for i, s in enumerate(sigs): - if i in timings: - continue - diff = sum(1 for (_, a), (_, b) in zip(best_sig, s) if a != b) - if diff != 1: - continue - if bench_idx(i) < best_ms: - improved = True - if not improved: - break + for _ in range(self.max_refine_iters): + best_idx = min(timings, key=timings.get) + best_sig = sigs[best_idx] + best_ms = timings[best_idx] + improved = False + for i, s in enumerate(sigs): + if i in timings: + continue + diff = sum(1 for (_, a), (_, b) in zip(best_sig, s) if a != b) + if diff != 1: + continue + if bench_idx(i) < best_ms: + improved = True + if not improved: + break # Crown RUNOFF: the search's single-shot timings carry enough noise to flip crowns # between same-code tune rolls (gpt-oss prefill grouped rolled 1727-2150us on identical @@ -761,7 +758,7 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): # without n_trials a 3-trial exploratory winner gets replayed at 100 trials. str(self.keys), f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," - f"basin_probes={self.basin_probes},refine={self.refine},runoff={self.runoff}", + f"basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters},runoff={self.runoff}", str(tuning_key), ] + [str(c) for c in configs] # str(Config) is pre_hook-free and process-stable cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest() @@ -898,7 +895,6 @@ def bayesian_autotune( n_startup_trials: int = 12, gamma: float = 0.25, basin_probes: bool = True, - refine: bool = True, runoff: bool = True, max_refine_iters: int = 5, log_path: str | None = None, @@ -915,10 +911,11 @@ def bayesian_autotune( n_startup_trials: random seed configs before the TPE model kicks in gamma: top fraction of measured configs treated as "good" basin_probes: measure categorical basin anchors before the TPE - refine, max_refine_iters: coordinate-descent refinement after the TPE + max_refine_iters: coordinate-descent iterations after the TPE (0 disables; env + override FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS) runoff: remeasure the top candidates before choosing the winner - (basin_probes/refine/runoff each have an env override: - FINEGRAINED_AUTOTUNE_BASIN_PROBES/_REFINE/_RUNOFF, 0 or 1) + (basin_probes/runoff each have an env override: + FINEGRAINED_AUTOTUNE_BASIN_PROBES/_RUNOFF, 0 or 1) path_anchor_axes: branch axes (kwarg names) whose values relocate the tile optimum — one guaranteed max-tile anchor per value combo (see _basin_anchor_indices); the kernel declares them, the @@ -945,7 +942,6 @@ def decorator(fn): n_startup_trials=n_startup_trials, gamma=gamma, basin_probes=basin_probes, - refine=refine, runoff=runoff, max_refine_iters=max_refine_iters, log_path=log_path, From 1f3fddb82f456bf4d3cddb9cfb233279a5236628 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 07:13:21 -0700 Subject: [PATCH 13/25] Choose autotuner runoff winner from fresh measurements --- finegrained-kernels/tests/test_autotuner.py | 17 +++++++++++++++++ .../finegrained_kernels/bayesian_autotuner.py | 11 +++++++---- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 377ee5284..f965c40eb 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -375,3 +375,20 @@ def kernel( assert [(c.kwargs["BLOCK_SIZE_M"], c.kwargs["BLOCK_SIZE_N"]) for c in anchors] == [ (64, 128), (128, 256), (64, 128), (128, 256), ] + + +@pytest.mark.kernels_ci +def test_runoff_winner_uses_remeasured_candidates(monkeypatch): + """A fourth tile's stale search sample must not beat the runoff's three fresh samples.""" + configs = [triton.Config({"BLOCK": block}) for block in (16, 32, 64, 128)] + + @bayesian_autotune(configs, [], n_trials=4, basin_probes=False, max_refine_iters=0) + @triton.jit + def kernel(BLOCK: tl.constexpr): + pass + + samples = {16: iter((1.0, 6.0)), 32: iter((2.0, 5.0)), + 64: iter((3.0, 7.0)), 128: iter((4.0,))} + monkeypatch.setattr(kernel, "_bench", lambda *args, config, **kwargs: next(samples[config.kwargs["BLOCK"]])) + winner = kernel._bayesian_search(configs, (), {}, ()) + assert winner.kwargs["BLOCK"] == 32 diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 8ec0a0e9d..9f274b54e 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -602,21 +602,24 @@ def bench_idx(idx: int) -> float: # the top candidates back-to-back — same warmed process, adjacent in time, so they see # the same clocks — and crown the winner of THAT comparison. Costs ~3 extra benches. finite = sorted((t, i) for i, t in timings.items() if t != float("inf")) + finalists = list(timings) if self.runoff and len(finite) > 1: - for _, i in finite[:3]: + finalists = [i for _, i in finite[:3]] + for i in finalists: timings.pop(i, None) bench_idx(i) # memoizing closure: re-benches and re-records the popped entry self.configs_timings = {configs[i]: t for i, t in timings.items()} - best = min(timings, key=timings.get) + # Other candidates still carry search-time samples, not comparable runoff measurements. + best = min(finalists, key=timings.get) if timings[best] == float("inf"): # crowning an arbitrary broken config would persist it as `best` on disk and # every later healthy process would replay it without re-tuning — the exact # sticky-contaminated-context scenario inf-forgiveness exists to survive self._report_bench_failures() raise RuntimeError( - f"[bayesian-autotune] {self.fn_name}: every benched config failed " - "(all-inf tune) — nothing correct to crown or persist; the context may " + f"[bayesian-autotune] {self.fn_name}: every final candidate failed " + "— nothing correct to crown or persist; the context may " "be trap-contaminated (retry in a fresh process) or the grid is broken " "for this launch (see the failure report above)." ) From 973a00529cd9cf84f3bd550049071cdc8d87af6a Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 07:19:45 -0700 Subject: [PATCH 14/25] Keep runoff regression independent of correctness tuning budget --- finegrained-kernels/tests/test_autotuner.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index f965c40eb..0597a8ab1 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -387,6 +387,9 @@ def test_runoff_winner_uses_remeasured_candidates(monkeypatch): def kernel(BLOCK: tl.constexpr): pass + # Exercise runoff even when the surrounding correctness run limits tuning to one trial. + monkeypatch.setattr(kernel, "n_trials", len(configs)) + monkeypatch.setattr(kernel, "runoff", True) samples = {16: iter((1.0, 6.0)), 32: iter((2.0, 5.0)), 64: iter((3.0, 7.0)), 128: iter((4.0,))} monkeypatch.setattr(kernel, "_bench", lambda *args, config, **kwargs: next(samples[config.kwargs["BLOCK"]])) From 37842c45494af2a58ea32a612f120ee9193186f7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 07:38:51 -0700 Subject: [PATCH 15/25] Bump kernel version for tuner signature change --- finegrained-kernels/build.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/finegrained-kernels/build.toml b/finegrained-kernels/build.toml index ee9751249..97fa5f78b 100644 --- a/finegrained-kernels/build.toml +++ b/finegrained-kernels/build.toml @@ -1,6 +1,6 @@ [general] name = "finegrained-kernels" -version = 0 +version = 1 license = "Apache-2.0" backends = [ "cuda", From b70b838df064605ed9f8cef3d17c9bf00c6dbf5c Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 15:13:29 -0700 Subject: [PATCH 16/25] Allow coupled configuration axes in Bayesian tuning --- finegrained-kernels/tests/test_autotuner.py | 29 +++++++++++++++++++ .../finegrained_kernels/bayesian_autotuner.py | 29 +++++++++++++++++-- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 0597a8ab1..6771e1d0b 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -395,3 +395,32 @@ def kernel(BLOCK: tl.constexpr): monkeypatch.setattr(kernel, "_bench", lambda *args, config, **kwargs: next(samples[config.kwargs["BLOCK"]])) winner = kernel._bayesian_search(configs, (), {}, ()) assert winner.kwargs["BLOCK"] == 32 + + +@pytest.mark.kernels_ci +def test_coupled_axes_cross_refinement_valley(monkeypatch): + """K and pipeline stages must change together when either change alone is slower.""" + configs = [ + triton.Config({"BLOCK_SIZE_K": bk}, num_stages=stages) + for bk in (128, 256) for stages in (2, 3) + ] + scores = {(128, 2): 10.0, (128, 3): 12.0, (256, 2): 12.0, (256, 3): 1.0} + winners = [] + for coupled_axes in ((), (("BLOCK_SIZE_K", "num_stages"),)): + @bayesian_autotune(configs, [], coupled_axes=coupled_axes, cache_results=False) + @triton.jit + def kernel(BLOCK_SIZE_K: tl.constexpr): + pass + + monkeypatch.setattr(kernel, "n_trials", 2) + monkeypatch.setattr(kernel, "n_startup_trials", 2) + monkeypatch.setattr(kernel, "basin_probes", True) + monkeypatch.setattr(kernel, "max_refine_iters", 1) + monkeypatch.setattr(kernel, "runoff", False) + monkeypatch.setattr(kernel, "_basin_anchor_indices", lambda configs: [0, 1]) + monkeypatch.setattr( + kernel, "_bench", lambda *args, config, **kwargs: scores[(config.kwargs["BLOCK_SIZE_K"], config.num_stages)] + ) + winner = kernel._bayesian_search(configs, (), {}, ()) + winners.append((winner.kwargs["BLOCK_SIZE_K"], winner.num_stages)) + assert winners == [(128, 2), (256, 3)] diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 9f274b54e..ff88876e8 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -93,6 +93,7 @@ def __init__( max_refine_iters: int = 5, log_path: str | None = None, path_anchor_axes: tuple[str, ...] = (), + coupled_axes: tuple[tuple[str, ...], ...] = (), finite_check_args: tuple[str, ...] = (), **kwargs, ): @@ -123,6 +124,12 @@ def __init__( # axis names). Each distinct value combination of these axes gets one guaranteed # max-tile anchor; empty = median anchors only. self.path_anchor_axes = tuple(path_anchor_axes) + # Joint choices share one TPE density and one coordinate-refinement step. + self.coupled_axes = tuple(tuple(group) for group in coupled_axes) + coupled_names = [name for group in self.coupled_axes for name in group] + if any(len(group) < 2 for group in self.coupled_axes) or len(set(coupled_names)) != len(coupled_names): + raise ValueError("coupled_axes must contain disjoint groups of at least two configuration fields") + self._coupled_names = frozenset(coupled_names) # Output arg NAMES (a declaration — the tuner never hardcodes kernel details) whose # values are checked for NaN/inf after each config's bench: a config that writes # non-finite output scores inf and can never be crowned. The tuner times, it never @@ -466,9 +473,22 @@ def benchmark(): self.nargs = None return ret + def _search_signatures(self, configs: List[Config]): + signatures = [] + for config in configs: + values = config.all_kwargs() + signature = {(name,): value for name, value in values.items() if name not in self._coupled_names} + for group in self.coupled_axes: + try: + signature[group] = tuple(values[name] for name in group) + except KeyError as exc: + raise ValueError(f"Unknown coupled_axes configuration field: {exc.args[0]}") from exc + signatures.append(tuple(sorted(signature.items()))) + return signatures + def _bayesian_search(self, configs: List[Config], args, kwargs, key) -> Config: timings: Dict[int, float] = {} - sigs = [tuple(sorted(c.all_kwargs().items())) for c in configs] + sigs = self._search_signatures(configs) def bench_idx(idx: int) -> float: if idx in timings: @@ -761,7 +781,8 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): # without n_trials a 3-trial exploratory winner gets replayed at 100 trials. str(self.keys), f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," - f"basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters},runoff={self.runoff}", + f"basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters},runoff={self.runoff}," + f"coupled_axes={self.coupled_axes}", str(tuning_key), ] + [str(c) for c in configs] # str(Config) is pre_hook-free and process-stable cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest() @@ -902,6 +923,7 @@ def bayesian_autotune( max_refine_iters: int = 5, log_path: str | None = None, path_anchor_axes: tuple[str, ...] = (), + coupled_axes: tuple[tuple[str, ...], ...] = (), finite_check_args: tuple[str, ...] = (), cache_results: bool = True, reset_to_zero=None, @@ -923,6 +945,8 @@ def bayesian_autotune( optimum — one guaranteed max-tile anchor per value combo (see _basin_anchor_indices); the kernel declares them, the tuner stays independent of configuration details + coupled_axes: groups of config fields treated as joint choices by TPE and + refinement, including Config fields such as num_stages max_failures: abort the tune once this many configs fail to compile/run before the measured budget is met (default n_trials; env override FINEGRAINED_AUTOTUNE_MAX_FAILURES). Without it a broadly broken @@ -949,6 +973,7 @@ def decorator(fn): max_refine_iters=max_refine_iters, log_path=log_path, path_anchor_axes=path_anchor_axes, + coupled_axes=coupled_axes, finite_check_args=finite_check_args, cache_results=cache_results, **kwargs, From f2ea548f30c595367746286b146a9052b450e1e0 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 15:21:51 -0700 Subject: [PATCH 17/25] Keep finegrained kernels on version zero --- finegrained-kernels/build.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/finegrained-kernels/build.toml b/finegrained-kernels/build.toml index 97fa5f78b..ee9751249 100644 --- a/finegrained-kernels/build.toml +++ b/finegrained-kernels/build.toml @@ -1,6 +1,6 @@ [general] name = "finegrained-kernels" -version = 1 +version = 0 license = "Apache-2.0" backends = [ "cuda", From bb49417efac31ea1a2a782b948a7c5e1a2e8f1da Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 15:30:02 -0700 Subject: [PATCH 18/25] Compare cold CUDA tuning costs with graph replay --- .../finegrained_kernels/bayesian_autotuner.py | 35 ++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index ff88876e8..a4ee3f4e9 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -37,9 +37,11 @@ import time from collections import defaultdict from contextlib import contextmanager +from functools import cached_property from pathlib import Path import torch +from triton.testing import do_bench_cudagraph from typing import Dict, List @@ -59,6 +61,25 @@ # Include the tuner implementation alongside Triton's kernel-source hash. _TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + +def _bench_cold_graph(fn, quantiles): + """Compare cold-input graph costs without per-launch timing events. + + Scores include the same cache-clearing work for every candidate; they are + comparative costs, not standalone kernel latencies. Timing a graph of repeated + flush/launch pairs amortizes events that can obscure short-kernel differences. + """ + cache = driver.active.get_empty_cache_for_benchmark() + + def call(): + driver.active.clear_cache(cache) + fn() + + # Triton's graph timer warms up on its own stream; finish input producers first. + torch.cuda.current_stream().synchronize() + return do_bench_cudagraph(call, quantiles=quantiles) + + # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every # BayesianAutotuner.run() call increments its kernel's tally, dumped to at exit. # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token @@ -146,6 +167,18 @@ def __init__( self._dtype_tag_cache: Dict[tuple, str] = {} self._dtype_str_cache: Dict[torch.dtype, str] = {} + @cached_property + def do_bench(self): + if self._do_bench is not None: + return self._do_bench + if ( + driver.active.get_current_target().backend == "cuda" + and not self.user_defined_pre_hook + and not self.user_defined_post_hook + ): + return _bench_cold_graph + return driver.active.get_benchmarker() + def _dtype_str(self, dtype) -> str: s = self._dtype_str_cache.get(dtype) if s is None: @@ -647,7 +680,7 @@ def bench_idx(idx: int) -> float: def _log_result(self, key, config: Config, ms: float): """Append one ``(key, config, ms)`` record as JSONL for offline analysis — - e.g. pruning configs that are consistently far off the per-key best. ``inf`` ms + CUDA graph costs include cache clearing, not just the kernel. ``inf`` ms marks a config that failed to compile/run (out of resources, etc.).""" if not self.log_path: return From e6410ae5c1edc134b52aa65f013c35122b96c8e7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 01:20:10 -0700 Subject: [PATCH 19/25] Order tuner phase options and document each environment override --- finegrained-kernels/README.md | 9 +++++---- finegrained-kernels/tests/test_autotuner.py | 2 +- .../finegrained_kernels/bayesian_autotuner.py | 16 ++++++++-------- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index 7398d8e9b..3e75e15eb 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,10 +157,11 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. Basin probes and winner remeasurement -are controlled by `FINEGRAINED_AUTOTUNE_BASIN_PROBES` and `FINEGRAINED_AUTOTUNE_RUNOFF` -(`0` or `1`, both enabled by default). `FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS` controls the -coordinate-refinement iteration limit (default `5`; `0` disables refinement). +`FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. `FINEGRAINED_AUTOTUNE_RUNOFF` +controls winner remeasurement (`0` or `1`, enabled by default). +`FINEGRAINED_AUTOTUNE_BASIN_PROBES` controls basin probes (`0` or `1`, enabled by default). +`FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS` controls the coordinate-refinement iteration limit +(default `5`; `0` disables refinement). For correctness tests, set the trial budget to `1` and all three controls to `0`: only one successful configuration is measured, retrying failed candidates. `FINEGRAINED_AUTOTUNE_LOG=` appends per-config timings as JSONL to ``. diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 6771e1d0b..78ae06ba6 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -414,9 +414,9 @@ def kernel(BLOCK_SIZE_K: tl.constexpr): monkeypatch.setattr(kernel, "n_trials", 2) monkeypatch.setattr(kernel, "n_startup_trials", 2) + monkeypatch.setattr(kernel, "runoff", False) monkeypatch.setattr(kernel, "basin_probes", True) monkeypatch.setattr(kernel, "max_refine_iters", 1) - monkeypatch.setattr(kernel, "runoff", False) monkeypatch.setattr(kernel, "_basin_anchor_indices", lambda configs: [0, 1]) monkeypatch.setattr( kernel, "_bench", lambda *args, config, **kwargs: scores[(config.kwargs["BLOCK_SIZE_K"], config.num_stages)] diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index a4ee3f4e9..e155bcaa6 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -109,8 +109,8 @@ def __init__( max_failures: int | None = None, n_startup_trials: int = 12, gamma: float = 0.25, - basin_probes: bool = True, runoff: bool = True, + basin_probes: bool = True, max_refine_iters: int = 5, log_path: str | None = None, path_anchor_axes: tuple[str, ...] = (), @@ -137,8 +137,8 @@ def __init__( self.n_startup_trials = n_startup_trials # top fraction of measured configs the TPE treats as "good" self.gamma = gamma - self.basin_probes = os.environ.get("FINEGRAINED_AUTOTUNE_BASIN_PROBES", str(int(basin_probes))) == "1" self.runoff = os.environ.get("FINEGRAINED_AUTOTUNE_RUNOFF", str(int(runoff))) == "1" + self.basin_probes = os.environ.get("FINEGRAINED_AUTOTUNE_BASIN_PROBES", str(int(basin_probes))) == "1" self.max_refine_iters = int(os.environ.get("FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS", str(max_refine_iters))) # Branch axes (by kwarg name) that relocate the TILE optimum — the caller's # declaration, since it is configuration knowledge (the tuner never hardcodes @@ -814,7 +814,7 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): # without n_trials a 3-trial exploratory winner gets replayed at 100 trials. str(self.keys), f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," - f"basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters},runoff={self.runoff}," + f"runoff={self.runoff},basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters}," f"coupled_axes={self.coupled_axes}", str(tuning_key), ] + [str(c) for c in configs] # str(Config) is pre_hook-free and process-stable @@ -951,8 +951,8 @@ def bayesian_autotune( max_failures: int | None = None, n_startup_trials: int = 12, gamma: float = 0.25, - basin_probes: bool = True, runoff: bool = True, + basin_probes: bool = True, max_refine_iters: int = 5, log_path: str | None = None, path_anchor_axes: tuple[str, ...] = (), @@ -968,12 +968,12 @@ def bayesian_autotune( that fail to compile/run are skipped without consuming it) n_startup_trials: random seed configs before the TPE model kicks in gamma: top fraction of measured configs treated as "good" + runoff: remeasure the top candidates before choosing the winner + (env override FINEGRAINED_AUTOTUNE_RUNOFF, 0 or 1) basin_probes: measure categorical basin anchors before the TPE + (env override FINEGRAINED_AUTOTUNE_BASIN_PROBES, 0 or 1) max_refine_iters: coordinate-descent iterations after the TPE (0 disables; env override FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS) - runoff: remeasure the top candidates before choosing the winner - (basin_probes/runoff each have an env override: - FINEGRAINED_AUTOTUNE_BASIN_PROBES/_RUNOFF, 0 or 1) path_anchor_axes: branch axes (kwarg names) whose values relocate the tile optimum — one guaranteed max-tile anchor per value combo (see _basin_anchor_indices); the kernel declares them, the @@ -1001,8 +1001,8 @@ def decorator(fn): max_failures=max_failures, n_startup_trials=n_startup_trials, gamma=gamma, - basin_probes=basin_probes, runoff=runoff, + basin_probes=basin_probes, max_refine_iters=max_refine_iters, log_path=log_path, path_anchor_axes=path_anchor_axes, From 73049e8db97a49ab6b9e3f6491934a66d6c73893 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 01:45:22 -0700 Subject: [PATCH 20/25] Keep Triton benchmark defaults in the Bayesian tuner --- .../finegrained_kernels/bayesian_autotuner.py | 34 +------------------ 1 file changed, 1 insertion(+), 33 deletions(-) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index e155bcaa6..439c0e1dc 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -37,11 +37,9 @@ import time from collections import defaultdict from contextlib import contextmanager -from functools import cached_property from pathlib import Path import torch -from triton.testing import do_bench_cudagraph from typing import Dict, List @@ -62,24 +60,6 @@ _TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() -def _bench_cold_graph(fn, quantiles): - """Compare cold-input graph costs without per-launch timing events. - - Scores include the same cache-clearing work for every candidate; they are - comparative costs, not standalone kernel latencies. Timing a graph of repeated - flush/launch pairs amortizes events that can obscure short-kernel differences. - """ - cache = driver.active.get_empty_cache_for_benchmark() - - def call(): - driver.active.clear_cache(cache) - fn() - - # Triton's graph timer warms up on its own stream; finish input producers first. - torch.cuda.current_stream().synchronize() - return do_bench_cudagraph(call, quantiles=quantiles) - - # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every # BayesianAutotuner.run() call increments its kernel's tally, dumped to at exit. # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token @@ -167,18 +147,6 @@ def __init__( self._dtype_tag_cache: Dict[tuple, str] = {} self._dtype_str_cache: Dict[torch.dtype, str] = {} - @cached_property - def do_bench(self): - if self._do_bench is not None: - return self._do_bench - if ( - driver.active.get_current_target().backend == "cuda" - and not self.user_defined_pre_hook - and not self.user_defined_post_hook - ): - return _bench_cold_graph - return driver.active.get_benchmarker() - def _dtype_str(self, dtype) -> str: s = self._dtype_str_cache.get(dtype) if s is None: @@ -680,7 +648,7 @@ def bench_idx(idx: int) -> float: def _log_result(self, key, config: Config, ms: float): """Append one ``(key, config, ms)`` record as JSONL for offline analysis — - CUDA graph costs include cache clearing, not just the kernel. ``inf`` ms + e.g. pruning configs that are consistently far off the per-key best. ``inf`` ms marks a config that failed to compile/run (out of resources, etc.).""" if not self.log_path: return From 431248a2ab7ab7fb302fb5eb0f5f2af500215d5d Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 09:53:55 -0700 Subject: [PATCH 21/25] Use cold-cache CUDA graph tuning with warm-cache opt-in --- finegrained-kernels/README.md | 5 +++ .../finegrained_kernels/bayesian_autotuner.py | 37 ++++++++++++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index 3e75e15eb..12199eddb 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,6 +157,11 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. +CUDA tuning uses graph timing of repeated cache-clear/kernel pairs to represent model weights +larger than L2; these comparative scores include cache clearing. Set +`FINEGRAINED_AUTOTUNE_WARM_CACHE=1` before importing the kernel to time repeated warm inputs +instead. The two modes have separate on-disk tuning results. +Explicit benchmark callbacks, legacy timing options and user hooks retain their existing behavior. `FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. `FINEGRAINED_AUTOTUNE_RUNOFF` controls winner remeasurement (`0` or `1`, enabled by default). `FINEGRAINED_AUTOTUNE_BASIN_PROBES` controls basin probes (`0` or `1`, enabled by default). diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 439c0e1dc..127604ffc 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -37,9 +37,11 @@ import time from collections import defaultdict from contextlib import contextmanager +from functools import cached_property, partial from pathlib import Path import torch +from triton.testing import do_bench_cudagraph from typing import Dict, List @@ -60,6 +62,25 @@ _TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() +def _bench_graph(fn, quantiles, *, warm_cache=False): + """Compare graph costs with cold inputs unless warm-cache timing is requested. + + Cold timing repeats cache-clear/kernel pairs while amortizing host/event overhead. + Those scores include the common cache-clearing cost; they are comparative costs, + not standalone kernel latencies. + """ + cache = None if warm_cache else driver.active.get_empty_cache_for_benchmark() + + def call(): + if cache is not None: + driver.active.clear_cache(cache) + fn() + + # Triton's graph timer warms up on its own stream; finish input producers first. + torch.cuda.current_stream().synchronize() + return do_bench_cudagraph(call, quantiles=quantiles) + + # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every # BayesianAutotuner.run() call increments its kernel's tally, dumped to at exit. # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token @@ -100,6 +121,8 @@ def __init__( ): super().__init__(*args, **kwargs) self.fn_name = getattr(self.fn, "__name__", str(self.fn)) + # Large-model weights exceed L2; opt into repeated warm inputs for smaller workloads. + self.warm_cache = os.environ.get("FINEGRAINED_AUTOTUNE_WARM_CACHE", "0") == "1" # Bayesian trial budget — the per-decorator default, overridable via the # FINEGRAINED_AUTOTUNE_TRIALS env var (quick sweeps / exhaustive runs without touching # the decorators; set it >= grid size to fall back to stock exhaustive bench-all). @@ -147,6 +170,18 @@ def __init__( self._dtype_tag_cache: Dict[tuple, str] = {} self._dtype_str_cache: Dict[torch.dtype, str] = {} + @cached_property + def do_bench(self): + if self._do_bench is not None: + return self._do_bench + if ( + driver.active.get_current_target().backend != "cuda" + or self.user_defined_pre_hook or self.user_defined_post_hook + ): + return driver.active.get_benchmarker() + + return partial(_bench_graph, warm_cache=self.warm_cache) + def _dtype_str(self, dtype) -> str: s = self._dtype_str_cache.get(dtype) if s is None: @@ -783,7 +818,7 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): str(self.keys), f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," f"runoff={self.runoff},basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters}," - f"coupled_axes={self.coupled_axes}", + f"coupled_axes={self.coupled_axes},warm_cache={self.warm_cache}", str(tuning_key), ] + [str(c) for c in configs] # str(Config) is pre_hook-free and process-stable cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest() From 66a93254654f8c41eb473ef0378742ffd2afcaca Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 10:53:50 -0700 Subject: [PATCH 22/25] Benchmark CUDA forwards with rotating cold weight banks --- finegrained-kernels/bench/README.md | 17 ++ finegrained-kernels/bench/bench_moe.py | 257 +++++++++++++++++-------- 2 files changed, 194 insertions(+), 80 deletions(-) diff --git a/finegrained-kernels/bench/README.md b/finegrained-kernels/bench/README.md index 52218a2f5..44a7fcc2a 100644 --- a/finegrained-kernels/bench/README.md +++ b/finegrained-kernels/bench/README.md @@ -30,6 +30,23 @@ Every cell runs in three modes: `eager`, `cudagraph` (decode's deployment mode), `compile` (`torch.compile(max-autotune, fullgraph)`). A cell that raises is a red ✕ marker; the others still run. +On CUDA, every mode cycles through identical weight banks whose **active** weights exceed +**twice the device's L2 cache**. For MoE, active bytes count only the experts selected by the +unchanged routing, including weight scales. Dense layers count the whole matrix. Each bank +has independent copies of the same inputs and routes, so compilation cannot share activation +quantization or routing across invocations. The timed callback returns every bank's output; +reported latency is its elapsed time divided by the bank count. Reuse inside one forward is +preserved. Prepared weight views and scales keep static addresses so compiled replay does not +warm them with staging copies. No cache-flush kernel or profiler time is included in graph latency. + +CSV rows record `cache_policy`, `bank_count`, total `active_weight_bytes` across the banks, +and `prepared_bytes` (the live allocation increase while the arm's bank closures are prepared; +for dense rows, this covers the shared set of arm closures). Logs also show total checkpoint +storage. Insufficient memory fails explicitly; it never falls back to warm weights. Replotting +or merging old CSVs without matching cache metadata requires rerunning the benchmark. XPU +retains its existing backend timing behavior and is labeled `backend_default`, not CUDA cold. +Previously generated figures used warm CUDA graph timing; regenerate them for cold-weight comparisons. + ## The pre-swizzled fast path By default the finegrained-kernels arm feeds **pre-swizzled** (`SWIZZLE_32_4_4`) MX weight scales, diff --git a/finegrained-kernels/bench/bench_moe.py b/finegrained-kernels/bench/bench_moe.py index 3b7099154..97e39ac34 100644 --- a/finegrained-kernels/bench/bench_moe.py +++ b/finegrained-kernels/bench/bench_moe.py @@ -52,6 +52,9 @@ cudagraph do_bench_cudagraph (decode's deployment mode) compile torch.compile(mode="max-autotune", fullgraph=True), warmed, then do_bench +CUDA modes rotate identical weight/input banks whose active weights exceed twice L2. +Elapsed latency is normalized per invocation; cache eviction is never timed. + Regimes: decode T=1 and prefill T=8192 (routed through top_k experts; attn row: M=T). --smoke: fast everything-compiles pass — 3-trial tunes (via FINEGRAINED_AUTOTUNE_TRIALS, set @@ -187,6 +190,7 @@ def _nvfp4_input_globals(cfg, hidden, gu, gus, gu_g): def amax(t): return (t.abs().amax() / (6.0 * 448.0)).clamp(min=1e-30).float().reshape(1) _NVFP4_GLOBALS[key] = (amax(hidden), amax(inter)) + _mark_static(*_NVFP4_GLOBALS[key]) return _NVFP4_GLOBALS[key] @@ -581,6 +585,7 @@ def _check_cuda_home(): ), } MODES = ["eager", "cudagraph", "compile"] +CACHE_POLICY = "rotating_gt_2l2" if DEVICE == "cuda" else "backend_default" IMPL_COLORS = { "finegrained-kernels": "#1f77b4", @@ -632,6 +637,7 @@ def _mark_static(*tensors): # a real model holds weights as Parameters (cudagraph-static); closure # tensors would be re-copied into the cudagraph buffers EVERY compiled # call (~150us/GB DtoD, measured 95% of compiled decode) — mark static. + # Captured views need their own mark; it is not inherited from their base. torch._dynamo.mark_static_address(t) @@ -683,6 +689,7 @@ def __init__(self, cfg, gu, dn, gus=None, dns=None): # the reference forwards read this like `act_fn`: a name when the model normalizes each # expert application before the routing weights, None when it does not self.post_expert_norm_name = None + self.has_post_expert_norm = False self._is_expert_parallel = False self.gate_up_proj, self.down_proj = gu, dn self.gate_up_proj_scale_inv, self.down_proj_scale_inv = gus, dns @@ -805,6 +812,7 @@ def torch_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, *_): gate_up_proj_weight_global_scale=gu_g, down_proj_weight_global_scale=dn_g, **nvfp4_kw) gus_b, dns_b = _torch_preblock_weight_scale(gus), _torch_preblock_weight_scale(dns) + _mark_static(gu, gus_b, dns_b) return lambda: fgm.moe.moe_torch_grouped(hidden, idx, w, gu, dn, gus_b, dns_b, **kw) @@ -999,6 +1007,8 @@ def swap(t): scale_c1 = scale_gate1 = scale_c2 = ones alpha, beta, limit = _trtllm_gated_scalars(cfg, E, hidden.device) act_format = _activation_format(cfg) + _mark_static(g1, g1s, d1, d1s, hidden_g, scale_c1, scale_gate1, scale_c2, + alpha, beta, limit) def run(): x = torch.nn.functional.pad(hidden, (0, Hp - H)) if Hp != H else hidden @@ -1091,6 +1101,7 @@ def trtllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g=None, dn if cfg["weights"] == "mxfp8": g1, g1s, dn_s, d1s = _trtllm_mxfp8_prep(gu, gus, dn, dns) alpha, beta, limit = _trtllm_gated_scalars(cfg, E, hidden.device) + _mark_static(g1, g1s, dn_s, d1s, alpha, beta, limit) def run(): hq, hsr = _fi_mxfp8_quantize(hidden, False) @@ -1111,6 +1122,7 @@ def run(): gus_f = gus.float() # UE8M0 first: cat has no e8m0 kernel g1s = torch.cat([gus_f[:, inter // 128:], gus_f[:, : inter // 128]], dim=1).contiguous() dns_f = dns.float().contiguous() + _mark_static(g1, g1s, dns_f) def run(): # the STACK's own act quant (vLLM's TRT-LLM integration runs this exact kernel) — @@ -1155,6 +1167,8 @@ def vllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): a2_scale=a["down_proj_activation_scale"].max().reshape(1).float(), per_act_token_quant=False, block_shape=None, ) + if qc is not None: + _mark_static(qc.w1_scale, qc.w2_scale, qc.a1_scale, qc.a2_scale) return lambda: vllm_fused_experts(hidden, gu, dn, tkw, tki, activation=_VllmAct.SILU, quant_config=qc) @@ -1188,6 +1202,7 @@ def prep(bf16_w): # (E, out, in) bf16 -> swizzled mxfp4 weight + precision conf tw, ws = tfmx.quantize_to_mxfp4( bf16_w.transpose(-1, -2).contiguous(), triton_kernels_hub) tw, ws = tfmx.swizzle_mxfp4(tw, ws, triton_kernels_hub) + _mark_static(tw.storage.data, ws.storage.data) return tw, pc(weight_scale=ws, flex_ctx=flex(rhs_data=inflex())) experts = tfmx.Mxfp4GptOssExperts( @@ -1264,6 +1279,7 @@ def megablocks_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): f"megablocks {name} wants {tuple(buf.shape)} ({buf.numel()} elems), " f"source has {tuple(src.shape)} ({src.numel()})") buf.copy_(src.contiguous().reshape(buf.shape).to(torch.bfloat16)) + _mark_static(buf) scores = torch.zeros(hidden.shape[0], E, device=hidden.device, dtype=torch.bfloat16) scores.scatter_(1, idx.long(), w.to(torch.bfloat16)) # belt and braces: some builds still hand back (out, bias) @@ -1313,27 +1329,67 @@ def _context_poisoned(tag, mode): return True -def bench_modes(run, tag): +def _tensor_bytes(tensors): + return sum(t.numel() * t.element_size() for t in tensors if t is not None) + + +def _clone_tensors(tensors): + return tuple(None if t is None else t.clone() for t in tensors) + + +def _weight_banks(weights, active_experts=None): + """Identical checkpoints with a cyclic active working set larger than CUDA L2.""" + weight_bytes = _tensor_bytes(weights) + active_bytes = (weight_bytes if active_experts is None + else weight_bytes * active_experts // weights[0].shape[0]) + count = (2 * ACCELERATOR.get_device_properties(0).L2_cache_size // active_bytes + 1 + if DEVICE == "cuda" else 1) + if DEVICE == "cuda": + free, _ = ACCELERATOR.mem_get_info() + free += ACCELERATOR.memory_reserved() - ACCELERATOR.memory_allocated() + required = (count - 1) * weight_bytes + if required > free: + raise RuntimeError(f"cold weight banks need {required / 1024**3:.2f} GiB extra, " + f"but only {free / 1024**3:.2f} GiB is free") + banks = [weights] + [_clone_tensors(weights) for _ in range(count - 1)] + for bank in banks: + _mark_static(*bank) + print(f" cache: {CACHE_POLICY}; banks={count}, " + f"active weights={count * active_bytes / 1024**2:.1f} MiB, " + f"allocated weights={count * weight_bytes / 1024**3:.2f} GiB", flush=True) + return banks, count * active_bytes + + +def bench_modes(runs, tag, active_weight_bytes, prepared_bytes): """{mode: latency_us | None}, plus the eager output for parity. A mode that raises is a red CRASH cell — the other modes still run (fresh error printed inline).""" - res, out = {}, None + def run(): + # Independent inputs prevent cross-bank CSE; all outputs must survive compilation. + return tuple(fn() for fn in runs) + + count = len(runs) + res = dict(bank_count=count, active_weight_bytes=active_weight_bytes, + prepared_bytes=prepared_bytes) + out = None try: out = run() torch.accelerator.synchronize() # warm + tune before ANY timing/capture - res["eager"] = do_bench(run, return_mode="min") * 1e3 + res["eager"] = do_bench(run, return_mode="min") * 1e3 / count print(f" {tag:14s} eager {res['eager']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} eager crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["eager"] = None _context_poisoned(tag, "eager") try: - res["cudagraph"] = do_bench_cudagraph(run, return_mode="min") * 1e3 + res["cudagraph"] = do_bench_cudagraph(run, return_mode="min") * 1e3 / count print(f" {tag:14s} cudagraph {res['cudagraph']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} cudagraph crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["cudagraph"] = None _context_poisoned(tag, "cudagraph") try: + # This wrapper is shared across every arm and regime; isolate its Dynamo guards. + torch._dynamo.reset_code(run.__code__) crun = torch.compile(run, mode="max-autotune", fullgraph=True) cout = crun() torch.accelerator.synchronize() @@ -1345,17 +1401,20 @@ def bench_modes(run, tag): # quantized path, and W4A8 spreads further than the cross-impl parity we already accept # on it (deepgemm-vs-ours 2.3e-2, its own compile-vs-eager 2.1e-2 — a 1e-2 bound flagged # the latter as a crash and put a false red X on the figure). - if isinstance(cout, torch.Tensor) and isinstance(out, torch.Tensor): - drift = rel_diff(cout, out) - if drift > 0.25: - raise RuntimeError(f"compiled output diverges from eager (rel {drift:.2e})") - res["compile"] = do_bench(crun, return_mode="min") * 1e3 + if len(cout) != count: + raise RuntimeError("compiled graph dropped a weight bank output") + for compiled, eager in zip(cout, out): + if isinstance(compiled, torch.Tensor) and isinstance(eager, torch.Tensor): + drift = rel_diff(compiled, eager) + if drift > 0.25: + raise RuntimeError(f"compiled output diverges from eager (rel {drift:.2e})") + res["compile"] = do_bench(crun, return_mode="min") * 1e3 / count print(f" {tag:14s} compile {res['compile']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} compile crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["compile"] = None _context_poisoned(tag, "compile") - return res, out + return res, None if out is None else out[0] def rel_diff(a, b): @@ -1387,23 +1446,28 @@ def bench_problem_row(row, pname, cfg, arms, weights, rows_out): # state so shapes don't go automatic-dynamic across problems and the shared # frames can't hit the fullgraph recompile limit torch._dynamo.reset() - gu, gus, dn, dns, gu_g, dn_g = weights for regime, tokens, grouped in (("decode", DECODE_TOKENS, False), ("prefill", PREFILL_TOKENS, True)): print(f" -- {regime}") hidden, idx, w, logits = routing(cfg, tokens) - args = (cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits) + banks, active_bytes = _weight_banks(weights, idx.unique().numel()) + inputs = [(hidden, idx, w, logits)] + [ + _clone_tensors((hidden, idx, w, logits)) for _ in banks[1:] + ] anchor_res, anchor_out = None, None for name in arms: try: - run = ARMS[name](*args) + before = ACCELERATOR.memory_allocated() + runs = [ARMS[name](cfg, grouped, h, i, rw, *bank, rl) + for bank, (h, i, rw, rl) in zip(banks, inputs)] + prepared_bytes = max(0, ACCELERATOR.memory_allocated() - before) except Exception as e: print(f" [{name} setup failed: {type(e).__name__}: {str(e)[:90]}]", flush=True) rows_out.append((row, pname, regime, _impl(name), {m: None for m in MODES}, None)) continue - res, out = bench_modes(run, name) + res, out = bench_modes(runs, name, active_bytes, prepared_bytes) parity = None if anchor_res is None: anchor_res, anchor_out = res, out @@ -1414,6 +1478,8 @@ def bench_problem_row(row, pname, cfg, arms, weights, rows_out): print(f" {name:14s} parity-vs-finegrained-kernels {parity:.1e}" f" finegrained-kernels-speedup {sp}", flush=True) rows_out.append((row, pname, regime, _impl(name), res, parity)) + del runs + del banks, inputs print() @@ -1422,27 +1488,9 @@ def _impl(arm_name): return arm_name.replace("_unfused", "").replace("_bf16", "") -def bench_attn_row(row, pname, cfg, rows_out): - """One attn linear per model, in its deployment format (same weights across - impls; the finegrained-kernels arm's ``activation_format`` follows the model — GPT-OSS runs W4A4).""" - print(f"== [{row}] {pname}") - if MOCK: - _mock_rows(row, pname, ("finegrained-kernels",) + cfg["baselines"], rows_out) - return - torch._dynamo.reset() +def _attn_arms(cfg, x, weights): N, K, block = cfg["N"], cfg["K"], cfg["block"] - W_g = None - if cfg["weights"] == "fp8_128x128_ue8m0": - W, Ws = make_weights(N, K, DEVICE, [128, 128], - scale_dtype=torch.float8_e8m0fnu) - else: - # registry makers are expert-batched; build E=1 and index the slab off. NVFP4 returns a - # per-expert fp32 global as its 3rd value (None for single-level FP8/MX) — keep it so the - # two-level global multiply isn't dropped (it was, silently running NVFP4 at global=1). - W, Ws, W_g = WEIGHTS[cfg["weights"]]["make"](N, K, 1) - W, Ws = W[0], Ws[0] - W_g = W_g if W_g is None else W_g[0] - _mark_static(W, Ws) + W, Ws, W_g = weights Ws_fp8 = _fp8_scales(Ws, block) dg_block = tuple(block) if block else None # deployment layout for the LOCAL arm only (every baseline reads the raw row-major scale): @@ -1452,6 +1500,7 @@ def bench_attn_row(row, pname, cfg, rows_out): if (PRESWIZZLE and DEVICE == "cuda" and cfg["weights"] in _MX_WEIGHTS and _quantized_acts(cfg) and N % 128 == 0): Ws_fgm = fgm.swizzle_mx_scales(Ws) + _mark_static(Ws_fgm) # OpenAI triton_kernels dense mxfp4 matmul (matmul_ogs, no routing): the qkv linear # in the GPT-OSS MXFP4 format. Weight is a single (1, K, N) expert, swizzled once # at load (same as the fused arm); latency-only (its own weights). @@ -1459,6 +1508,7 @@ def bench_attn_row(row, pname, cfg, rows_out): tw_bf = torch.randn(1, K, N, device=DEVICE, dtype=torch.bfloat16) * 0.05 tw, tws = _tfmx.quantize_to_mxfp4(tw_bf, triton_kernels_hub) tw, tws = _tfmx.swizzle_mxfp4(tw, tws, triton_kernels_hub) + _mark_static(tw.storage.data, tws.storage.data) tk_pc = triton_kernels_hub.matmul_ogs.PrecisionConfig( weight_scale=tws, flex_ctx=triton_kernels_hub.matmul_ogs.FlexCtx( @@ -1489,6 +1539,7 @@ def bench_attn_row(row, pname, cfg, rows_out): rb = [ScalingType.BlockWise1x16, ScalingType.TensorWise] swz = [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE] mat_b = W.view(FP4).t() + _mark_static(mat_b, *sb, one) def torch_mm(a, a_scale=None): aq, a_s = fgm.nvfp4_act_quant(a) @@ -1501,6 +1552,7 @@ def torch_mm(a, a_scale=None): rb = ScalingType.BlockWise1x32 swz = SwizzleType.SWIZZLE_32_4_4 mat_b = W.t() + _mark_static(mat_b, sb) def torch_mm(a, a_scale=None): aq, a_s = fgm.mxfp8_act_quant(a) @@ -1516,6 +1568,7 @@ def torch_mm(a, a_scale=None): from torch.nn.functional import ScalingType mat_b = W.t() + _mark_static(mat_b) def torch_mm(a, a_scale): return torch.nn.functional.scaled_mm( @@ -1524,53 +1577,86 @@ def torch_mm(a, a_scale): Ws.reshape(()), ScalingType.TensorWise, output_dtype=torch.bfloat16) + # a STATIC row's calibrated activation scale replaces the inline quant: one value, since a + # dense module is one quantized module (the MoE's per-expert form has no dense analogue) + attn_act_scale = ( + (x.float().abs().amax() / 448.0).clamp(min=1e-12).reshape(1) + if cfg.get("static") else None + ) + if attn_act_scale is not None: + _mark_static(attn_act_scale) + # act is inline-quantized (As=None) unless the row is static; Ws is the weight scale (Bs); activation_format is the + # activation precision (None follows the weight format). + attn_arms = { + "finegrained-kernels": lambda: fgm.matmul_2d( + x, W, attn_act_scale, Ws_fgm, activation_format=_activation_format(cfg), output_dtype=torch.bfloat16, + b_global_scale=W_g), + } + if torch_mm is not None: + attn_arms["torch_mm"] = lambda: torch_mm(x, attn_act_scale) + if "finegrained-fp8" in cfg["baselines"]: + attn_arms["finegrained-fp8"] = lambda: upstream_fp8.matmul_2d(x, W, Ws_fp8, block, + torch.bfloat16) + if "deepgemm" in cfg["baselines"]: + attn_arms["deepgemm"] = lambda: deepgemm_fp8_fp4_linear( + x, W, Ws, block_size=dg_block, output_dtype=torch.bfloat16) + if "triton_kernels" in cfg["baselines"]: + attn_arms["triton_kernels"] = lambda: tk_ogs( + x, tw, None, None, precision_config=tk_pc) + # The dense GEMM is where transformers' NVFP4 path actually lives, so this is the + # apples-to-apples cell: same weight, its packer, its kernel. Packing is offline (the + # checkpoint conversion does it), so only `gemm` is timed. + if _nvfp4_gemm is not None and cfg["weights"] == "nvfp4": + # their packer takes a bf16 weight and produces its OWN layout, so feed it the + # dequantized tensor — `W` is already packed E2M1 (uint8, two values per byte), and + # casting those bytes to bf16 hands it a K/2-wide matrix ("expected 3072 input + # features, got 6144"). Same underlying weight, each stack quantizing it its own way, + # which is how every other baseline here is treated. + _w_bf16 = WEIGHTS["nvfp4"]["dequant"]( + W[None], Ws[None], (W_g if W_g is not None else torch.ones(1, device=W.device)) + )[0].to(torch.bfloat16) + _pw = _nvfp4_gemm.pack(_w_bf16, device=x.device) + _mark_static(_pw.qweight, _pw.sf, _pw.global_scale, _pw.sf_rowmajor) + attn_arms["nvfp4_gemm"] = lambda: _nvfp4_gemm.gemm(_pw, x) + return attn_arms + + +def bench_attn_row(row, pname, cfg, rows_out): + """One attn linear per model, in its deployment format (same weights across + impls; the finegrained-kernels arm's ``activation_format`` follows the model — GPT-OSS runs W4A4).""" + print(f"== [{row}] {pname}") + if MOCK: + _mock_rows(row, pname, ("finegrained-kernels",) + cfg["baselines"], rows_out) + return + torch._dynamo.reset() + N, K = cfg["N"], cfg["K"] + W_g = None + if cfg["weights"] == "fp8_128x128_ue8m0": + W, Ws = make_weights(N, K, DEVICE, [128, 128], + scale_dtype=torch.float8_e8m0fnu) + else: + # registry makers are expert-batched; build E=1 and index the slab off. NVFP4 returns a + # per-expert fp32 global as its 3rd value (None for single-level FP8/MX) — keep it so the + # two-level global multiply isn't dropped (it was, silently running NVFP4 at global=1). + W, Ws, W_g = WEIGHTS[cfg["weights"]]["make"](N, K, 1) + W, Ws = W[0], Ws[0] + W_g = W_g if W_g is None else W_g[0] + _mark_static(W, Ws) + weights = (W, Ws, W_g) + for regime, tokens in (("decode", DECODE_TOKENS), ("prefill", PREFILL_TOKENS)): print(f" -- {regime}") torch.manual_seed(0) x = torch.randn(tokens, K, device=DEVICE, dtype=torch.bfloat16) - # a STATIC row's calibrated activation scale replaces the inline quant: one value, since a - # dense module is one quantized module (the MoE's per-expert form has no dense analogue) - attn_act_scale = ( - (x.float().abs().amax() / 448.0).clamp(min=1e-12).reshape(1) - if cfg.get("static") else None - ) - if attn_act_scale is not None: - _mark_static(attn_act_scale) - # act is inline-quantized (As=None) unless the row is static; Ws is the weight scale (Bs); activation_format is the - # activation precision (None follows the weight format). - attn_arms = { - "finegrained-kernels": lambda: fgm.matmul_2d( - x, W, attn_act_scale, Ws_fgm, activation_format=_activation_format(cfg), output_dtype=torch.bfloat16, - b_global_scale=W_g), - } - if torch_mm is not None: - attn_arms["torch_mm"] = lambda: torch_mm(x, attn_act_scale) - if "finegrained-fp8" in cfg["baselines"]: - attn_arms["finegrained-fp8"] = lambda: upstream_fp8.matmul_2d(x, W, Ws_fp8, block, - torch.bfloat16) - if "deepgemm" in cfg["baselines"]: - attn_arms["deepgemm"] = lambda: deepgemm_fp8_fp4_linear( - x, W, Ws, block_size=dg_block, output_dtype=torch.bfloat16) - if "triton_kernels" in cfg["baselines"]: - attn_arms["triton_kernels"] = lambda: tk_ogs( - x, tw, None, None, precision_config=tk_pc) - # The dense GEMM is where transformers' NVFP4 path actually lives, so this is the - # apples-to-apples cell: same weight, its packer, its kernel. Packing is offline (the - # checkpoint conversion does it), so only `gemm` is timed. - if _nvfp4_gemm is not None and cfg["weights"] == "nvfp4": - # their packer takes a bf16 weight and produces its OWN layout, so feed it the - # dequantized tensor — `W` is already packed E2M1 (uint8, two values per byte), and - # casting those bytes to bf16 hands it a K/2-wide matrix ("expected 3072 input - # features, got 6144"). Same underlying weight, each stack quantizing it its own way, - # which is how every other baseline here is treated. - _w_bf16 = WEIGHTS["nvfp4"]["dequant"]( - W[None], Ws[None], (W_g if W_g is not None else torch.ones(1, device=W.device)) - )[0].to(torch.bfloat16) - _pw = _nvfp4_gemm.pack(_w_bf16, device=x.device) - attn_arms["nvfp4_gemm"] = lambda: _nvfp4_gemm.gemm(_pw, x) + banks, active_bytes = _weight_banks(weights) + inputs = [x] + [x.clone() for _ in banks[1:]] + before = ACCELERATOR.memory_allocated() + bank_arms = [_attn_arms(cfg, hidden, bank) for bank, hidden in zip(banks, inputs)] + prepared_bytes = max(0, ACCELERATOR.memory_allocated() - before) anchor_res, anchor_out = None, None - for name, run in attn_arms.items(): - res, out = bench_modes(run, name) + for name in bank_arms[0]: + runs = [arms[name] for arms in bank_arms] + res, out = bench_modes(runs, name, active_bytes, prepared_bytes) parity = None if anchor_res is None: anchor_res, anchor_out = res, out @@ -1581,11 +1667,12 @@ def torch_mm(a, a_scale): print(f" {name:14s} parity-vs-finegrained-kernels {parity:.1e}" f" finegrained-kernels-speedup {sp}", flush=True) rows_out.append((row, pname, regime, name, res, parity)) + del runs, bank_arms, banks, inputs print() device_name = "MOCK (random values)" if MOCK else ACCELERATOR.get_device_name(0) -print(f"device: {device_name} torch {torch.__version__}" +print(f"device: {device_name} torch {torch.__version__} cache={CACHE_POLICY}" f"{' [SMOKE]' if SMOKE else ''}") print("finegrained-kernels = local build; baselines: finegrained-fp8 (upstream), DeepGEMM, " "vLLM fused MoE, FlashInfer's TRT-LLM routed MoE, transformers grouped_mm/batched_mm, " @@ -1601,7 +1688,8 @@ def wanted(*names): # owning one device via CUDA_VISIBLE_DEVICES). The coordinator spawns GPUS workers, each # writes a shard CSV, then the coordinator merges + plots. A single GPU (GPUS=1) runs inline. ── _CSV = os.path.join(_HERE, "bench_moe.csv") -_CSV_HEADER = "category,problem,regime,impl,mode,latency_us,parity_vs_finegrained_kernels\n" +_CSV_HEADER = ("category,problem,regime,impl,mode,latency_us,parity_vs_finegrained_kernels," + "cache_policy,bank_count,active_weight_bytes,prepared_bytes\n") def _write_rows_csv(path, rows_out): @@ -1612,7 +1700,9 @@ def _write_rows_csv(path, rows_out): v = res.get(mode) f.write(f'"{rr}","{p}",{reg},{i},{mode},' f'{"" if v is None else f"{v:.2f}"},' - f'{"" if par is None else par}\n') + f'{"" if par is None else par},{CACHE_POLICY},' + f'{res.get("bank_count", "")},{res.get("active_weight_bytes", "")},' + f'{res.get("prepared_bytes", "")}\n') def _load_rows_csv(path): @@ -1638,11 +1728,17 @@ def _allowed(cfg): allowed["linear quantized", pn] = _allowed(c) | {"torch_mm", "nvfp4_gemm"} acc = {} # (cat, problem, regime, impl) -> (res dict, parity) for r in csv.DictReader(open(path)): + if r.get("cache_policy") != CACHE_POLICY: + raise ValueError(f"{path}: missing/mismatched cache policy; rerun the benchmark " + "instead of mixing old warm-graph results") if r["impl"] not in allowed.get((r["category"], r["problem"]), {"finegrained-kernels"}): continue key = (r["category"], r["problem"], r["regime"], r["impl"]) res, par = acc.setdefault(key, ({}, None)) res[r["mode"]] = float(r["latency_us"]) if r["latency_us"] else None + for field in ("bank_count", "active_weight_bytes", "prepared_bytes"): + if r.get(field): + res[field] = int(r[field]) if r["parity_vs_finegrained_kernels"]: acc[key] = (res, float(r["parity_vs_finegrained_kernels"])) return [(cat, p, reg, impl, res, par) @@ -1936,9 +2032,10 @@ def _drawn(res): handles.append(mpatches.Patch(facecolor="none", edgecolor="black", hatch="..", label="eager (no cudagraph)")) ax.legend(handles=handles, loc="upper left", fontsize=8) +cache_note = "cold weights (>2xL2)" if DEVICE == "cuda" else "backend cache defaults" fig.suptitle(f"MoE bench — finegrained-kernels vs finegrained-fp8 + references " f"({device_name}, real model shapes; decode=cudagraph+compile, " - f"prefill=eager)", y=0.9995) + f"prefill=eager; {cache_note})", y=0.9995) fig.tight_layout(rect=(0, 0, 1, 0.99)) out_png = os.path.join(_HERE, f"bench_moe{suffix}.png") fig.savefig(out_png, dpi=120) From 12cc0d0b63eb28de09dd919294e17b1157b06615 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 11:06:09 -0700 Subject: [PATCH 23/25] Share timing helpers and exclude eviction from autotuner scores --- finegrained-kernels/README.md | 7 ++-- finegrained-kernels/bench/bench_moe.py | 4 +-- .../finegrained_kernels/bayesian_autotuner.py | 30 ++++------------- .../torch-ext/finegrained_kernels/bench.py | 33 +++++++++++++++++++ 4 files changed, 45 insertions(+), 29 deletions(-) create mode 100644 finegrained-kernels/torch-ext/finegrained_kernels/bench.py diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index 12199eddb..73722cfef 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,10 +157,9 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -CUDA tuning uses graph timing of repeated cache-clear/kernel pairs to represent model weights -larger than L2; these comparative scores include cache clearing. Set -`FINEGRAINED_AUTOTUNE_WARM_CACHE=1` before importing the kernel to time repeated warm inputs -instead. The two modes have separate on-disk tuning results. +CUDA tuning clears L2 before each measured event interval; cache clearing is excluded +from the score. Set `FINEGRAINED_AUTOTUNE_WARM_CACHE=1` before importing the kernel to +opt into warm CUDA graph timing instead. The two modes have separate on-disk tuning results. Explicit benchmark callbacks, legacy timing options and user hooks retain their existing behavior. `FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. `FINEGRAINED_AUTOTUNE_RUNOFF` controls winner remeasurement (`0` or `1`, enabled by default). diff --git a/finegrained-kernels/bench/bench_moe.py b/finegrained-kernels/bench/bench_moe.py index 97e39ac34..03f52c2af 100644 --- a/finegrained-kernels/bench/bench_moe.py +++ b/finegrained-kernels/bench/bench_moe.py @@ -124,13 +124,13 @@ os.environ.setdefault("FINEGRAINED_AUTOTUNE_TRIALS", "3") import torch # noqa: E402 -from triton.testing import do_bench, do_bench_cudagraph # noqa: E402 _HERE = os.path.dirname(os.path.abspath(__file__)) _ROOT = os.path.dirname(_HERE) sys.path.insert(0, os.path.join(_ROOT, "torch-ext")) sys.path.insert(0, os.path.join(_ROOT, "tests")) import finegrained_kernels as fgm # noqa: E402 local branch +from finegrained_kernels.bench import do_bench # noqa: E402 from kernels import get_kernel # noqa: E402 # kernels-community IS a trusted publisher, so the check passes whenever it can reach the Hub @@ -1381,7 +1381,7 @@ def run(): res["eager"] = None _context_poisoned(tag, "eager") try: - res["cudagraph"] = do_bench_cudagraph(run, return_mode="min") * 1e3 / count + res["cudagraph"] = do_bench(run, cudagraph=True, return_mode="min") * 1e3 / count print(f" {tag:14s} cudagraph {res['cudagraph']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} cudagraph crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 127604ffc..369e89457 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -41,7 +41,6 @@ from pathlib import Path import torch -from triton.testing import do_bench_cudagraph from typing import Dict, List @@ -56,29 +55,14 @@ triton_key, ) -logger = logging.getLogger(__name__) - -# Include the tuner implementation alongside Triton's kernel-source hash. -_TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() - +from .bench import do_bench -def _bench_graph(fn, quantiles, *, warm_cache=False): - """Compare graph costs with cold inputs unless warm-cache timing is requested. - - Cold timing repeats cache-clear/kernel pairs while amortizing host/event overhead. - Those scores include the common cache-clearing cost; they are comparative costs, - not standalone kernel latencies. - """ - cache = None if warm_cache else driver.active.get_empty_cache_for_benchmark() - - def call(): - if cache is not None: - driver.active.clear_cache(cache) - fn() +logger = logging.getLogger(__name__) - # Triton's graph timer warms up on its own stream; finish input producers first. - torch.cuda.current_stream().synchronize() - return do_bench_cudagraph(call, quantiles=quantiles) +# Include timing and tuner implementations alongside Triton's kernel-source hash. +_TUNER_SOURCE_HASH = hashlib.sha256( + Path(__file__).read_bytes() + Path(__file__).with_name("bench.py").read_bytes() +).hexdigest() # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every @@ -180,7 +164,7 @@ def do_bench(self): ): return driver.active.get_benchmarker() - return partial(_bench_graph, warm_cache=self.warm_cache) + return partial(do_bench, cudagraph=self.warm_cache) def _dtype_str(self, dtype) -> str: s = self._dtype_str_cache.get(dtype) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bench.py b/finegrained-kernels/torch-ext/finegrained_kernels/bench.py new file mode 100644 index 000000000..202357869 --- /dev/null +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bench.py @@ -0,0 +1,33 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Timing shared by autotuning and benchmarks.""" + +import torch +from triton.testing import do_bench as _do_bench +from triton.testing import do_bench_cudagraph + + +def do_bench(fn, quantiles=None, *, cudagraph=False, **kwargs): + """Time execution in milliseconds, excluding cache eviction. + + Event timing clears L2 before each measured interval. Graph timing has no + eviction: cold-weight callers must rotate a working set larger than L2 inside + ``fn``; otherwise it measures warm replay. + """ + if not cudagraph: + return _do_bench(fn, quantiles=quantiles, **kwargs) + # Triton's graph timer warms up on its own stream; finish input producers first. + torch.cuda.current_stream().synchronize() + return do_bench_cudagraph(fn, quantiles=quantiles, **kwargs) From e9c48cf034b6131878e7109b5a7a41b5033a3dd1 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 11:35:03 -0700 Subject: [PATCH 24/25] Tune forward GEMMs with rotating cold weights in CUDA graphs --- finegrained-kernels/README.md | 11 +- finegrained-kernels/bench/bench_moe.py | 7 +- finegrained-kernels/tests/test_autotuner.py | 37 +++++++ .../torch-ext/finegrained_kernels/batched.py | 7 ++ .../finegrained_kernels/bayesian_autotuner.py | 74 +++++++++++-- .../torch-ext/finegrained_kernels/bench.py | 100 ++++++++++++++++++ .../torch-ext/finegrained_kernels/grouped.py | 9 +- .../torch-ext/finegrained_kernels/matmul.py | 7 ++ 8 files changed, 238 insertions(+), 14 deletions(-) diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index 73722cfef..b8e722541 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,9 +157,14 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -CUDA tuning clears L2 before each measured event interval; cache clearing is excluded -from the score. Set `FINEGRAINED_AUTOTUNE_WARM_CACHE=1` before importing the kernel to -opt into warm CUDA graph timing instead. The two modes have separate on-disk tuning results. +Forward GEMMs tune with CUDA graph replay over identical weight banks whose active +weight values exceed twice L2. Routed kernels count only selected expert planes. Weight +strides and host TMA bindings are preserved; activations, scales and outputs stay shared. +Banks are reused across candidates and released after the search. Unsupported layouts, +excessive bank counts or insufficient memory use cold event timing, with cache clearing +outside the measured interval. Kernels without a declared weight input also use cold events. +Set `FINEGRAINED_AUTOTUNE_WARM_CACHE=1` before importing the kernel to opt into warm graph +timing instead. The two modes have separate on-disk tuning results. Explicit benchmark callbacks, legacy timing options and user hooks retain their existing behavior. `FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. `FINEGRAINED_AUTOTUNE_RUNOFF` controls winner remeasurement (`0` or `1`, enabled by default). diff --git a/finegrained-kernels/bench/bench_moe.py b/finegrained-kernels/bench/bench_moe.py index 03f52c2af..b4ded1fb3 100644 --- a/finegrained-kernels/bench/bench_moe.py +++ b/finegrained-kernels/bench/bench_moe.py @@ -130,7 +130,7 @@ sys.path.insert(0, os.path.join(_ROOT, "torch-ext")) sys.path.insert(0, os.path.join(_ROOT, "tests")) import finegrained_kernels as fgm # noqa: E402 local branch -from finegrained_kernels.bench import do_bench # noqa: E402 +from finegrained_kernels.bench import do_bench, rotation_count # noqa: E402 from kernels import get_kernel # noqa: E402 # kernels-community IS a trusted publisher, so the check passes whenever it can reach the Hub @@ -1342,8 +1342,9 @@ def _weight_banks(weights, active_experts=None): weight_bytes = _tensor_bytes(weights) active_bytes = (weight_bytes if active_experts is None else weight_bytes * active_experts // weights[0].shape[0]) - count = (2 * ACCELERATOR.get_device_properties(0).L2_cache_size // active_bytes + 1 - if DEVICE == "cuda" else 1) + count = rotation_count(active_bytes, 0) if DEVICE == "cuda" else 1 + if count is None: + raise RuntimeError("cold weight working set needs too many rotation banks") if DEVICE == "cuda": free, _ = ACCELERATOR.mem_get_info() free += ACCELERATOR.memory_reserved() - ACCELERATOR.memory_allocated() diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 78ae06ba6..53630fea6 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -19,6 +19,7 @@ import logging import subprocess import sys +import weakref import pytest import torch @@ -30,6 +31,7 @@ import finegrained_kernels # type: ignore import finegrained_kernels.matmul # type: ignore from finegrained_kernels.bayesian_autotuner import bayesian_autotune # type: ignore +from finegrained_kernels.bench import weight_input # type: ignore from finegrained_kernels.compat import is_sm10x # type: ignore from finegrained_kernels.pruners import block_within_dim_pruner, mx_config_pruner # type: ignore @@ -244,6 +246,41 @@ def _budget_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): assert len(measured) == 3, _budget_kernel.configs_timings +@pytest.mark.kernels_ci +@pytest.mark.skipif(TEST_DEVICE != "cuda", reason="CUDA graph timing required") +@pytest.mark.parametrize("n_trials", [2, 3]) +def test_rotating_weights_are_reused_and_released(monkeypatch, n_trials): + """Cold graph tuning reuses identical banks across configs and frees them after the search.""" + weight = torch.randn(256, device=TEST_DEVICE) + output = torch.empty_like(weight) + monkeypatch.setattr("finegrained_kernels.bench.rotation_count", lambda active_bytes, device: 3) + pointers, references = {}, {} + + def observe(arguments): + value = arguments["B"] + pointers.setdefault(arguments["BLOCK"], set()).add(value.data_ptr()) + references[value.data_ptr()] = weakref.ref(value) + + @bayesian_autotune( + [triton.Config({"BLOCK": block}, pre_hook=observe) for block in (32, 64, 128)], + ["N"], + n_trials=n_trials, + benchmark_input=weight_input, + cache_results=False, + ) + @triton.jit + def copy_weight(B, C, N, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tl.store(C + offsets, tl.load(B + offsets, offsets < N, 0), offsets < N) + + copy_weight[lambda meta: (triton.cdiv(256, meta["BLOCK"]),)](weight, output, 256) + torch.testing.assert_close(output, weight) + assert len(pointers) >= 2 + assert len({frozenset(values) for values in pointers.values()}) == 1 + assert all(len(values) == 3 for values in pointers.values()) + assert all(reference() is None for pointer, reference in references.items() if pointer != weight.data_ptr()) + + @pytest.mark.kernels_ci @pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") def test_act_quant_arms_are_bit_equal(): diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/batched.py b/finegrained-kernels/torch-ext/finegrained_kernels/batched.py index 043e379fd..46cdee294 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/batched.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/batched.py @@ -22,6 +22,7 @@ from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait from .bayesian_autotuner import bayesian_autotune +from .bench import weight_input from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype, pdl_launch_kwargs from .descriptors import rebind_batched_mx_bs_descriptor @@ -104,6 +105,7 @@ def expert_setup( # two checkpoints at one shape with different blocks must not share a winner. ["N", "K", "S", "OUTPUT_FORMAT", "GATE", "BLOCK_N", "BLOCK_K"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ "early_config_prune": compose_pruners( @@ -248,6 +250,7 @@ def w8a8_block_dynamic_fp8_matmul_batched_kernel( # keyed like the block-dynamic sibling: requant narrows the legal tiles to the quant block ["N", "K", "S", "OUTPUT_FORMAT", "GATE", "BLOCK_N", "BLOCK_K"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ "early_config_prune": compose_pruners( @@ -382,6 +385,7 @@ def w8a8_block_static_fp8_matmul_batched_kernel( get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True), ["N", "K", "S"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is # row-masked only — veto non-dividing tiles on both. @@ -532,6 +536,7 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( # incidentally today; this kernel's requant Cs is always row-major — no SWIZZLED_OUT axis). ["N", "K", "S", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BK-within-K + the sm_10x MMA-shape guards (swapped dot_scaled needs BN >= 128 for the # native scaled-MMA; smaller-BN swap configs never win and mislead the TPE). @@ -727,6 +732,7 @@ def mx_dynamic_matmul_batched_kernel( ), ["N", "K", "S", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ # No BK-within-K veto: this kernel's tail tile slides back in bounds (see the K-loop), @@ -867,6 +873,7 @@ def mx_weight_only_matmul_batched_kernel( # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide). ["N", "K", "S", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is # row-masked only — veto non-dividing tiles on both. diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index 369e89457..f7958fa47 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -55,7 +55,7 @@ triton_key, ) -from .bench import do_bench +from .bench import do_bench, weight_banks logger = logging.getLogger(__name__) @@ -101,12 +101,15 @@ def __init__( path_anchor_axes: tuple[str, ...] = (), coupled_axes: tuple[tuple[str, ...], ...] = (), finite_check_args: tuple[str, ...] = (), + benchmark_input=None, **kwargs, ): super().__init__(*args, **kwargs) self.fn_name = getattr(self.fn, "__name__", str(self.fn)) # Large-model weights exceed L2; opt into repeated warm inputs for smaller workloads. self.warm_cache = os.environ.get("FINEGRAINED_AUTOTUNE_WARM_CACHE", "0") == "1" + self.benchmark_input = benchmark_input + self._benchmark_banks = None # Bayesian trial budget — the per-decorator default, overridable via the # FINEGRAINED_AUTOTUNE_TRIALS env var (quick sweeps / exhaustive runs without touching # the decorators; set it >= grid size to fall back to stock exhaustive bench-all). @@ -207,7 +210,15 @@ def _bench(self, *args, config, **meta): self._failures.append((config, memo[sig] + " [memoized]")) return [float("inf")] * 3 try: - timing = super()._bench(*args, config=config, **meta) + if ( + self.benchmark_input is not None and not self.warm_cache + and self._do_bench is None + and not self.user_defined_pre_hook and not self.user_defined_post_hook + and driver.active.get_current_target().backend == "cuda" + ): + timing = self._bench_rotating(*args, config=config, **meta) + else: + timing = super()._bench(*args, config=config, **meta) bad = self._nonfinite_output(meta) if bad is not None: self._failures.append((config, f"non-finite output in {bad} (numerics veto)")) @@ -336,6 +347,43 @@ def _persist_failed_compile_memo(self): except Exception: pass # persistence is best-effort; the in-memory memo still holds + def _bench_rotating(self, *args, config, **meta): + if self._benchmark_banks is None: + arguments = {**dict(zip(self.arg_names, args)), **meta} + weight, active_bytes = self.benchmark_input(arguments) + mutated = self.reset_to_zero + self.restore_value + if any(arguments.get(name) is weight for name in mutated): + self._benchmark_banks = False + else: + self._benchmark_banks = weight_banks(arguments, weight, active_bytes) or False + if self._benchmark_banks is False: + logger.info("[autotune] %s: weight rotation unavailable; using cold event timing", self.fn_name) + if self._benchmark_banks is False: + return super()._bench(*args, config=config, **meta) + + config_args = config.all_kwargs() + conflicts = meta.keys() & config.kwargs.keys() + if conflicts: + raise ValueError(f"Conflicting meta-parameters: {', '.join(conflicts)}") + calls = [{**bank, **config_args} for bank in self._benchmark_banks] + + def run(): + for arguments in calls: + if config.pre_hook is not None: + config.pre_hook(arguments) + self.pre_hook(arguments) + try: + self.fn.run(**arguments) + except Exception as exc: + try: + self.post_hook(arguments, exception=exc) + finally: + raise + self.post_hook(arguments, exception=None) + + timing = do_bench(run, quantiles=(0.5, 0.2, 0.8), cudagraph=True) + return [value / len(calls) for value in timing] + def _report_bench_failures(self): """After every tune, report every UNIQUE failure — a failure is never silent: inf-scoring keeps the search alive, but a human must see what broke (e.g. a code @@ -396,7 +444,10 @@ def run(self, *args, **kwargs): if config_kwargs is None: config_kwargs = config._all_kwargs_memo = config.all_kwargs() return self.fn.run(*args, **kwargs, **config_kwargs) - ret = super().run(*args, **kwargs) + try: + ret = super().run(*args, **kwargs) + finally: + self._benchmark_banks = None self._report_bench_failures() return ret @@ -461,10 +512,13 @@ def benchmark(): full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()} self.pre_hook(full_nargs, reset_only=True) - if self.cache_results: - self.check_disk_cache(key, pruned, benchmark) - else: - benchmark() + try: + if self.cache_results: + self.check_disk_cache(key, pruned, benchmark) + else: + benchmark() + finally: + self._benchmark_banks = None config = self.cache[key] self.best_config = config @@ -800,6 +854,8 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): # and FINEGRAINED_AUTOTUNE_TRIALS is not in Triton's invalidating-env list — # without n_trials a 3-trial exploratory winner gets replayed at 100 trials. str(self.keys), + str((getattr(self.benchmark_input, "__module__", None), + getattr(self.benchmark_input, "__qualname__", None))), f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," f"runoff={self.runoff},basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters}," f"coupled_axes={self.coupled_axes},warm_cache={self.warm_cache}", @@ -945,6 +1001,7 @@ def bayesian_autotune( path_anchor_axes: tuple[str, ...] = (), coupled_axes: tuple[tuple[str, ...], ...] = (), finite_check_args: tuple[str, ...] = (), + benchmark_input=None, cache_results: bool = True, reset_to_zero=None, restore_value=None, @@ -967,6 +1024,8 @@ def bayesian_autotune( tuner stays independent of configuration details coupled_axes: groups of config fields treated as joint choices by TPE and refinement, including Config fields such as num_stages + benchmark_input: return a read-only weight tensor and its active read bytes + from the launch arguments for cold rotating graph timing max_failures: abort the tune once this many configs fail to compile/run before the measured budget is met (default n_trials; env override FINEGRAINED_AUTOTUNE_MAX_FAILURES). Without it a broadly broken @@ -995,6 +1054,7 @@ def decorator(fn): path_anchor_axes=path_anchor_axes, coupled_axes=coupled_axes, finite_check_args=finite_check_args, + benchmark_input=benchmark_input, cache_results=cache_results, **kwargs, ) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bench.py b/finegrained-kernels/torch-ext/finegrained_kernels/bench.py index 202357869..9742c6ba1 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bench.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bench.py @@ -14,7 +14,10 @@ """Timing shared by autotuning and benchmarks.""" +from dataclasses import replace + import torch +from triton.tools.tensor_descriptor import TensorDescriptor from triton.testing import do_bench as _do_bench from triton.testing import do_bench_cudagraph @@ -31,3 +34,100 @@ def do_bench(fn, quantiles=None, *, cudagraph=False, **kwargs): # Triton's graph timer warms up on its own stream; finish input producers first. torch.cuda.current_stream().synchronize() return do_bench_cudagraph(fn, quantiles=quantiles, **kwargs) + + +def weight_input(nargs): + """Main GEMM weight and its logical read bytes, excluding inactive expert planes.""" + weight = nargs["B"] + active_bytes = weight.numel() * weight.element_size() + expert_ids = nargs.get("ExpertIds") + expert_start = nargs.get("ExpertStart") + if expert_ids is not None: + experts = weight.shape[0] + active = expert_ids[(expert_ids >= 0) & (expert_ids < experts)].unique().numel() + active_bytes = active_bytes * active // experts + elif expert_start is not None: + experts = weight.shape[0] + if expert_start.numel() < experts + 1: + return weight, 0 + active = (expert_start[1:experts + 1] > expert_start[:experts]).sum().item() + active_bytes = active_bytes * active // experts + return weight, active_bytes + + +def rotation_count(active_bytes, device): + """Cycle more than twice L2, or decline an impractically large captured graph.""" + if active_bytes <= 0: + return None + count = 2 * torch.cuda.get_device_properties(device).L2_cache_size // active_bytes + 1 + return count if count <= 1024 else None + + +def weight_banks(nargs, weight, active_bytes): + """Rotate readonly weight values; return None when cold event timing is safer. + + Activations, scales and outputs stay shared. Raw tensor aliases of the weight + are unsupported; host descriptors may alias views wholly inside its storage span. + """ + if not isinstance(weight, torch.Tensor) or not weight.is_cuda: + return None + if active_bytes > weight.numel() * weight.element_size(): + return None + # Broadcast/overlapping or holed layouts cannot count logical bytes as unique reads. + if torch._debug_has_internal_overlap(weight) != 0: + return None + names = [name for name, value in nargs.items() if value is weight] + if len(names) != 1: + return None + name = names[0] + storage = weight.untyped_storage()._cdata + descriptors = {} + span_bytes = weight.numel() * weight.element_size() + for key, value in nargs.items(): + if key != name and isinstance(value, torch.Tensor): + if value.untyped_storage()._cdata == storage: + return None + if isinstance(value, TensorDescriptor): + if not isinstance(value.base, torch.Tensor): + return None + if value.base.untyped_storage()._cdata != storage: + continue + base = value.base + offset = base.data_ptr() - weight.data_ptr() + extent = (1 + sum((size - 1) * stride for size, stride in zip(base.shape, base.stride()))) + if (offset < 0 or offset % 16 or offset % base.element_size() + or offset + extent * base.element_size() > span_bytes): + return None + descriptors[key] = (value, offset) + count = rotation_count(active_bytes, weight.device) + if count is None: + return None + free, _ = torch.cuda.mem_get_info(weight.device) + free += (torch.cuda.memory_reserved(weight.device) + - torch.cuda.memory_allocated(weight.device)) + if (count - 1) * span_bytes > free: + return None + banks = [nargs] + try: + with torch.no_grad(): + for _ in range(count - 1): + clone = torch.empty_strided(weight.shape, weight.stride(), + dtype=weight.dtype, device=weight.device) + clone.copy_(weight) + bank = dict(nargs, **{name: clone}) + rebound = {} + for key, (descriptor, offset) in descriptors.items(): + if id(descriptor) not in rebound: + base = descriptor.base + view = torch.empty(0, dtype=base.dtype, device=base.device).set_( + clone.untyped_storage(), offset // base.element_size(), + base.shape, base.stride(), + ) + rebound[id(descriptor)] = replace( + descriptor, base=view, block_shape=list(descriptor.block_shape), + ) + bank[key] = rebound[id(descriptor)] + banks.append(bank) + except torch.OutOfMemoryError: + return None + return banks diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py b/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py index 9423bdb28..55cf77d6c 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py @@ -21,7 +21,8 @@ from triton.tools.tensor_descriptor import TensorDescriptor from .bayesian_autotuner import bayesian_autotune -from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, persistent_program_count, sm_count, tl_dtype +from .bench import weight_input +from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, persistent_program_count, tl_dtype from .descriptors import build_grouped_operand_descriptors, rebind_grouped_descriptors, rebind_grouped_mx_descriptors from .formats import check_activation_format, global_scale_stride, is_per_expert_global, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, routed_rows, tokens_per_expert_bucket, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, mx_act_quant_grouped, mx_act_quant_routed_rows, quantize_routed_rows_per_expert, static_expert_act_operands, swizzle_grouped_mx_scales, tensor_wide_act_operands @@ -63,6 +64,7 @@ # descriptor-box legality, so checkpoints with different blocks must not share a winner. ["N", "K", "tokens_per_expert_bit_length", "GATE", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # Pipeliner-race guard: per launch-BM, WS-only at BM >= 64 and non-WS below (see the pruner). prune_configs_by={ @@ -267,6 +269,7 @@ def w8a8_block_dynamic_fp8_matmul_grouped_kernel( ), ["N", "K", "tokens_per_expert_bit_length", "GATE", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # Pipeliner-race WS guard (per launch-BM: forces WS where the default pipeliner races at # BM>=64, keeps non-WS below) + descriptor-box limits. @@ -457,6 +460,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel( # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is # row-masked only — veto non-dividing tiles on both. WS is a pure perf axis here @@ -661,6 +665,7 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( # splits them incidentally today; a descriptor Cs has no dtype token). ["N", "K", "tokens_per_expert_bit_length", "GATE", "SWIZZLED_SCALES", "OUTPUT_FORMAT", "SWIZZLED_OUT"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BK-within-K veto + the sm_10x dot_scaled shape/trap gates (this kernel had no # pruner while its BK span was {128,256} — the union span's BK=64 rows made the @@ -926,6 +931,7 @@ def mx_dynamic_matmul_grouped_kernel( ), ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ "early_config_prune": compose_pruners( @@ -1095,6 +1101,7 @@ def mx_weight_only_matmul_grouped_kernel( # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K is a tuned axis and the K-loop is maskless — veto non-dividing BKs. prune_configs_by={ diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py b/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py index 2d624fa7d..f44d3b80c 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py @@ -21,6 +21,7 @@ from triton.tools.tensor_descriptor import TensorDescriptor from .bayesian_autotuner import bayesian_autotune +from .bench import weight_input from .compat import add_op_namespace_prefix, FP8_DTYPE, is_sm10x, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype from .descriptors import maybe_descriptor, rebind_bd_descriptors, rebind_mx_descriptors, rebind_weight_only_descriptors from .formats import check_activation_format, global_scale_stride, normalize_global_scale, e2m1_as_uint8, is_mx, mx_scale_family, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_2d_operands, weight_format @@ -93,6 +94,7 @@ # at one shape with different blocks must not share a winner. ["N", "K", "m_bit_length", "GATE", "block_n", "block_k"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -283,6 +285,7 @@ def w8a8_block_dynamic_fp8_matmul_kernel( # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide, a distinct config space). ["N", "K", "m_bit_length", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -415,6 +418,7 @@ def w8a8_tensor_dynamic_fp8_matmul_kernel( # block_n/block_k (the launch-pinned quant block) gate descriptor-box legality and the K tile. ["N", "K", "m_bit_length", "GATE", "block_n", "block_k"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -588,6 +592,7 @@ def w8a8_block_static_fp8_matmul_kernel( # dtype token, so don't lean on the append). ["N", "K", "m_bit_length", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT", "SWIZZLED_OUT"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -787,6 +792,7 @@ def mx_dynamic_matmul_kernel( ), ["N", "K", "m_bit_length", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -927,6 +933,7 @@ def mx_weight_only_matmul_2d_kernel( # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide). ["N", "K", "m_bit_length", "GATE"], n_trials=100, + benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), From 6f19acef9c8a5cf2d00d3c9eff7ef95cc4a17459 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 18 Oct 2026 11:49:36 -0700 Subject: [PATCH 25/25] Separate cold-cache timing from the tuner search fixes --- finegrained-kernels/README.md | 9 - finegrained-kernels/bench/README.md | 17 -- finegrained-kernels/bench/bench_moe.py | 260 ++++++------------ finegrained-kernels/tests/test_autotuner.py | 37 --- .../torch-ext/finegrained_kernels/batched.py | 7 - .../finegrained_kernels/bayesian_autotuner.py | 97 +------ .../torch-ext/finegrained_kernels/bench.py | 133 --------- .../torch-ext/finegrained_kernels/grouped.py | 9 +- .../torch-ext/finegrained_kernels/matmul.py | 7 - 9 files changed, 91 insertions(+), 485 deletions(-) delete mode 100644 finegrained-kernels/torch-ext/finegrained_kernels/bench.py diff --git a/finegrained-kernels/README.md b/finegrained-kernels/README.md index b8e722541..3e75e15eb 100644 --- a/finegrained-kernels/README.md +++ b/finegrained-kernels/README.md @@ -157,15 +157,6 @@ the offline activation quants the ops use internally, exposed for pre-quantized Every kernel is tuned by a Tree-structured Parzen Estimator (TPE, the Bayesian search Optuna uses) autotuner with per-shape disk caching, config pruners that fence compiler bugs and can't-win regions per (arch, format), and failed-compile memoization. -Forward GEMMs tune with CUDA graph replay over identical weight banks whose active -weight values exceed twice L2. Routed kernels count only selected expert planes. Weight -strides and host TMA bindings are preserved; activations, scales and outputs stay shared. -Banks are reused across candidates and released after the search. Unsupported layouts, -excessive bank counts or insufficient memory use cold event timing, with cache clearing -outside the measured interval. Kernels without a declared weight input also use cold events. -Set `FINEGRAINED_AUTOTUNE_WARM_CACHE=1` before importing the kernel to opt into warm graph -timing instead. The two modes have separate on-disk tuning results. -Explicit benchmark callbacks, legacy timing options and user hooks retain their existing behavior. `FINEGRAINED_AUTOTUNE_TRIALS` overrides the trial budget. `FINEGRAINED_AUTOTUNE_RUNOFF` controls winner remeasurement (`0` or `1`, enabled by default). `FINEGRAINED_AUTOTUNE_BASIN_PROBES` controls basin probes (`0` or `1`, enabled by default). diff --git a/finegrained-kernels/bench/README.md b/finegrained-kernels/bench/README.md index 44a7fcc2a..52218a2f5 100644 --- a/finegrained-kernels/bench/README.md +++ b/finegrained-kernels/bench/README.md @@ -30,23 +30,6 @@ Every cell runs in three modes: `eager`, `cudagraph` (decode's deployment mode), `compile` (`torch.compile(max-autotune, fullgraph)`). A cell that raises is a red ✕ marker; the others still run. -On CUDA, every mode cycles through identical weight banks whose **active** weights exceed -**twice the device's L2 cache**. For MoE, active bytes count only the experts selected by the -unchanged routing, including weight scales. Dense layers count the whole matrix. Each bank -has independent copies of the same inputs and routes, so compilation cannot share activation -quantization or routing across invocations. The timed callback returns every bank's output; -reported latency is its elapsed time divided by the bank count. Reuse inside one forward is -preserved. Prepared weight views and scales keep static addresses so compiled replay does not -warm them with staging copies. No cache-flush kernel or profiler time is included in graph latency. - -CSV rows record `cache_policy`, `bank_count`, total `active_weight_bytes` across the banks, -and `prepared_bytes` (the live allocation increase while the arm's bank closures are prepared; -for dense rows, this covers the shared set of arm closures). Logs also show total checkpoint -storage. Insufficient memory fails explicitly; it never falls back to warm weights. Replotting -or merging old CSVs without matching cache metadata requires rerunning the benchmark. XPU -retains its existing backend timing behavior and is labeled `backend_default`, not CUDA cold. -Previously generated figures used warm CUDA graph timing; regenerate them for cold-weight comparisons. - ## The pre-swizzled fast path By default the finegrained-kernels arm feeds **pre-swizzled** (`SWIZZLE_32_4_4`) MX weight scales, diff --git a/finegrained-kernels/bench/bench_moe.py b/finegrained-kernels/bench/bench_moe.py index b4ded1fb3..3b7099154 100644 --- a/finegrained-kernels/bench/bench_moe.py +++ b/finegrained-kernels/bench/bench_moe.py @@ -52,9 +52,6 @@ cudagraph do_bench_cudagraph (decode's deployment mode) compile torch.compile(mode="max-autotune", fullgraph=True), warmed, then do_bench -CUDA modes rotate identical weight/input banks whose active weights exceed twice L2. -Elapsed latency is normalized per invocation; cache eviction is never timed. - Regimes: decode T=1 and prefill T=8192 (routed through top_k experts; attn row: M=T). --smoke: fast everything-compiles pass — 3-trial tunes (via FINEGRAINED_AUTOTUNE_TRIALS, set @@ -124,13 +121,13 @@ os.environ.setdefault("FINEGRAINED_AUTOTUNE_TRIALS", "3") import torch # noqa: E402 +from triton.testing import do_bench, do_bench_cudagraph # noqa: E402 _HERE = os.path.dirname(os.path.abspath(__file__)) _ROOT = os.path.dirname(_HERE) sys.path.insert(0, os.path.join(_ROOT, "torch-ext")) sys.path.insert(0, os.path.join(_ROOT, "tests")) import finegrained_kernels as fgm # noqa: E402 local branch -from finegrained_kernels.bench import do_bench, rotation_count # noqa: E402 from kernels import get_kernel # noqa: E402 # kernels-community IS a trusted publisher, so the check passes whenever it can reach the Hub @@ -190,7 +187,6 @@ def _nvfp4_input_globals(cfg, hidden, gu, gus, gu_g): def amax(t): return (t.abs().amax() / (6.0 * 448.0)).clamp(min=1e-30).float().reshape(1) _NVFP4_GLOBALS[key] = (amax(hidden), amax(inter)) - _mark_static(*_NVFP4_GLOBALS[key]) return _NVFP4_GLOBALS[key] @@ -585,7 +581,6 @@ def _check_cuda_home(): ), } MODES = ["eager", "cudagraph", "compile"] -CACHE_POLICY = "rotating_gt_2l2" if DEVICE == "cuda" else "backend_default" IMPL_COLORS = { "finegrained-kernels": "#1f77b4", @@ -637,7 +632,6 @@ def _mark_static(*tensors): # a real model holds weights as Parameters (cudagraph-static); closure # tensors would be re-copied into the cudagraph buffers EVERY compiled # call (~150us/GB DtoD, measured 95% of compiled decode) — mark static. - # Captured views need their own mark; it is not inherited from their base. torch._dynamo.mark_static_address(t) @@ -689,7 +683,6 @@ def __init__(self, cfg, gu, dn, gus=None, dns=None): # the reference forwards read this like `act_fn`: a name when the model normalizes each # expert application before the routing weights, None when it does not self.post_expert_norm_name = None - self.has_post_expert_norm = False self._is_expert_parallel = False self.gate_up_proj, self.down_proj = gu, dn self.gate_up_proj_scale_inv, self.down_proj_scale_inv = gus, dns @@ -812,7 +805,6 @@ def torch_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, *_): gate_up_proj_weight_global_scale=gu_g, down_proj_weight_global_scale=dn_g, **nvfp4_kw) gus_b, dns_b = _torch_preblock_weight_scale(gus), _torch_preblock_weight_scale(dns) - _mark_static(gu, gus_b, dns_b) return lambda: fgm.moe.moe_torch_grouped(hidden, idx, w, gu, dn, gus_b, dns_b, **kw) @@ -1007,8 +999,6 @@ def swap(t): scale_c1 = scale_gate1 = scale_c2 = ones alpha, beta, limit = _trtllm_gated_scalars(cfg, E, hidden.device) act_format = _activation_format(cfg) - _mark_static(g1, g1s, d1, d1s, hidden_g, scale_c1, scale_gate1, scale_c2, - alpha, beta, limit) def run(): x = torch.nn.functional.pad(hidden, (0, Hp - H)) if Hp != H else hidden @@ -1101,7 +1091,6 @@ def trtllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g=None, dn if cfg["weights"] == "mxfp8": g1, g1s, dn_s, d1s = _trtllm_mxfp8_prep(gu, gus, dn, dns) alpha, beta, limit = _trtllm_gated_scalars(cfg, E, hidden.device) - _mark_static(g1, g1s, dn_s, d1s, alpha, beta, limit) def run(): hq, hsr = _fi_mxfp8_quantize(hidden, False) @@ -1122,7 +1111,6 @@ def run(): gus_f = gus.float() # UE8M0 first: cat has no e8m0 kernel g1s = torch.cat([gus_f[:, inter // 128:], gus_f[:, : inter // 128]], dim=1).contiguous() dns_f = dns.float().contiguous() - _mark_static(g1, g1s, dns_f) def run(): # the STACK's own act quant (vLLM's TRT-LLM integration runs this exact kernel) — @@ -1167,8 +1155,6 @@ def vllm_moe_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): a2_scale=a["down_proj_activation_scale"].max().reshape(1).float(), per_act_token_quant=False, block_shape=None, ) - if qc is not None: - _mark_static(qc.w1_scale, qc.w2_scale, qc.a1_scale, qc.a2_scale) return lambda: vllm_fused_experts(hidden, gu, dn, tkw, tki, activation=_VllmAct.SILU, quant_config=qc) @@ -1202,7 +1188,6 @@ def prep(bf16_w): # (E, out, in) bf16 -> swizzled mxfp4 weight + precision conf tw, ws = tfmx.quantize_to_mxfp4( bf16_w.transpose(-1, -2).contiguous(), triton_kernels_hub) tw, ws = tfmx.swizzle_mxfp4(tw, ws, triton_kernels_hub) - _mark_static(tw.storage.data, ws.storage.data) return tw, pc(weight_scale=ws, flex_ctx=flex(rhs_data=inflex())) experts = tfmx.Mxfp4GptOssExperts( @@ -1279,7 +1264,6 @@ def megablocks_arm(cfg, grouped, hidden, idx, w, gu, gus, dn, dns, *_): f"megablocks {name} wants {tuple(buf.shape)} ({buf.numel()} elems), " f"source has {tuple(src.shape)} ({src.numel()})") buf.copy_(src.contiguous().reshape(buf.shape).to(torch.bfloat16)) - _mark_static(buf) scores = torch.zeros(hidden.shape[0], E, device=hidden.device, dtype=torch.bfloat16) scores.scatter_(1, idx.long(), w.to(torch.bfloat16)) # belt and braces: some builds still hand back (out, bias) @@ -1329,68 +1313,27 @@ def _context_poisoned(tag, mode): return True -def _tensor_bytes(tensors): - return sum(t.numel() * t.element_size() for t in tensors if t is not None) - - -def _clone_tensors(tensors): - return tuple(None if t is None else t.clone() for t in tensors) - - -def _weight_banks(weights, active_experts=None): - """Identical checkpoints with a cyclic active working set larger than CUDA L2.""" - weight_bytes = _tensor_bytes(weights) - active_bytes = (weight_bytes if active_experts is None - else weight_bytes * active_experts // weights[0].shape[0]) - count = rotation_count(active_bytes, 0) if DEVICE == "cuda" else 1 - if count is None: - raise RuntimeError("cold weight working set needs too many rotation banks") - if DEVICE == "cuda": - free, _ = ACCELERATOR.mem_get_info() - free += ACCELERATOR.memory_reserved() - ACCELERATOR.memory_allocated() - required = (count - 1) * weight_bytes - if required > free: - raise RuntimeError(f"cold weight banks need {required / 1024**3:.2f} GiB extra, " - f"but only {free / 1024**3:.2f} GiB is free") - banks = [weights] + [_clone_tensors(weights) for _ in range(count - 1)] - for bank in banks: - _mark_static(*bank) - print(f" cache: {CACHE_POLICY}; banks={count}, " - f"active weights={count * active_bytes / 1024**2:.1f} MiB, " - f"allocated weights={count * weight_bytes / 1024**3:.2f} GiB", flush=True) - return banks, count * active_bytes - - -def bench_modes(runs, tag, active_weight_bytes, prepared_bytes): +def bench_modes(run, tag): """{mode: latency_us | None}, plus the eager output for parity. A mode that raises is a red CRASH cell — the other modes still run (fresh error printed inline).""" - def run(): - # Independent inputs prevent cross-bank CSE; all outputs must survive compilation. - return tuple(fn() for fn in runs) - - count = len(runs) - res = dict(bank_count=count, active_weight_bytes=active_weight_bytes, - prepared_bytes=prepared_bytes) - out = None + res, out = {}, None try: out = run() torch.accelerator.synchronize() # warm + tune before ANY timing/capture - res["eager"] = do_bench(run, return_mode="min") * 1e3 / count + res["eager"] = do_bench(run, return_mode="min") * 1e3 print(f" {tag:14s} eager {res['eager']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} eager crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["eager"] = None _context_poisoned(tag, "eager") try: - res["cudagraph"] = do_bench(run, cudagraph=True, return_mode="min") * 1e3 / count + res["cudagraph"] = do_bench_cudagraph(run, return_mode="min") * 1e3 print(f" {tag:14s} cudagraph {res['cudagraph']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} cudagraph crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["cudagraph"] = None _context_poisoned(tag, "cudagraph") try: - # This wrapper is shared across every arm and regime; isolate its Dynamo guards. - torch._dynamo.reset_code(run.__code__) crun = torch.compile(run, mode="max-autotune", fullgraph=True) cout = crun() torch.accelerator.synchronize() @@ -1402,20 +1345,17 @@ def run(): # quantized path, and W4A8 spreads further than the cross-impl parity we already accept # on it (deepgemm-vs-ours 2.3e-2, its own compile-vs-eager 2.1e-2 — a 1e-2 bound flagged # the latter as a crash and put a false red X on the figure). - if len(cout) != count: - raise RuntimeError("compiled graph dropped a weight bank output") - for compiled, eager in zip(cout, out): - if isinstance(compiled, torch.Tensor) and isinstance(eager, torch.Tensor): - drift = rel_diff(compiled, eager) - if drift > 0.25: - raise RuntimeError(f"compiled output diverges from eager (rel {drift:.2e})") - res["compile"] = do_bench(crun, return_mode="min") * 1e3 / count + if isinstance(cout, torch.Tensor) and isinstance(out, torch.Tensor): + drift = rel_diff(cout, out) + if drift > 0.25: + raise RuntimeError(f"compiled output diverges from eager (rel {drift:.2e})") + res["compile"] = do_bench(crun, return_mode="min") * 1e3 print(f" {tag:14s} compile {res['compile']:9.1f}us", flush=True) except Exception as e: print(f" [{tag} compile crashed: {type(e).__name__}: {str(e)[:90]}]", flush=True) res["compile"] = None _context_poisoned(tag, "compile") - return res, None if out is None else out[0] + return res, out def rel_diff(a, b): @@ -1447,28 +1387,23 @@ def bench_problem_row(row, pname, cfg, arms, weights, rows_out): # state so shapes don't go automatic-dynamic across problems and the shared # frames can't hit the fullgraph recompile limit torch._dynamo.reset() + gu, gus, dn, dns, gu_g, dn_g = weights for regime, tokens, grouped in (("decode", DECODE_TOKENS, False), ("prefill", PREFILL_TOKENS, True)): print(f" -- {regime}") hidden, idx, w, logits = routing(cfg, tokens) - banks, active_bytes = _weight_banks(weights, idx.unique().numel()) - inputs = [(hidden, idx, w, logits)] + [ - _clone_tensors((hidden, idx, w, logits)) for _ in banks[1:] - ] + args = (cfg, grouped, hidden, idx, w, gu, gus, dn, dns, gu_g, dn_g, logits) anchor_res, anchor_out = None, None for name in arms: try: - before = ACCELERATOR.memory_allocated() - runs = [ARMS[name](cfg, grouped, h, i, rw, *bank, rl) - for bank, (h, i, rw, rl) in zip(banks, inputs)] - prepared_bytes = max(0, ACCELERATOR.memory_allocated() - before) + run = ARMS[name](*args) except Exception as e: print(f" [{name} setup failed: {type(e).__name__}: {str(e)[:90]}]", flush=True) rows_out.append((row, pname, regime, _impl(name), {m: None for m in MODES}, None)) continue - res, out = bench_modes(runs, name, active_bytes, prepared_bytes) + res, out = bench_modes(run, name) parity = None if anchor_res is None: anchor_res, anchor_out = res, out @@ -1479,8 +1414,6 @@ def bench_problem_row(row, pname, cfg, arms, weights, rows_out): print(f" {name:14s} parity-vs-finegrained-kernels {parity:.1e}" f" finegrained-kernels-speedup {sp}", flush=True) rows_out.append((row, pname, regime, _impl(name), res, parity)) - del runs - del banks, inputs print() @@ -1489,9 +1422,27 @@ def _impl(arm_name): return arm_name.replace("_unfused", "").replace("_bf16", "") -def _attn_arms(cfg, x, weights): +def bench_attn_row(row, pname, cfg, rows_out): + """One attn linear per model, in its deployment format (same weights across + impls; the finegrained-kernels arm's ``activation_format`` follows the model — GPT-OSS runs W4A4).""" + print(f"== [{row}] {pname}") + if MOCK: + _mock_rows(row, pname, ("finegrained-kernels",) + cfg["baselines"], rows_out) + return + torch._dynamo.reset() N, K, block = cfg["N"], cfg["K"], cfg["block"] - W, Ws, W_g = weights + W_g = None + if cfg["weights"] == "fp8_128x128_ue8m0": + W, Ws = make_weights(N, K, DEVICE, [128, 128], + scale_dtype=torch.float8_e8m0fnu) + else: + # registry makers are expert-batched; build E=1 and index the slab off. NVFP4 returns a + # per-expert fp32 global as its 3rd value (None for single-level FP8/MX) — keep it so the + # two-level global multiply isn't dropped (it was, silently running NVFP4 at global=1). + W, Ws, W_g = WEIGHTS[cfg["weights"]]["make"](N, K, 1) + W, Ws = W[0], Ws[0] + W_g = W_g if W_g is None else W_g[0] + _mark_static(W, Ws) Ws_fp8 = _fp8_scales(Ws, block) dg_block = tuple(block) if block else None # deployment layout for the LOCAL arm only (every baseline reads the raw row-major scale): @@ -1501,7 +1452,6 @@ def _attn_arms(cfg, x, weights): if (PRESWIZZLE and DEVICE == "cuda" and cfg["weights"] in _MX_WEIGHTS and _quantized_acts(cfg) and N % 128 == 0): Ws_fgm = fgm.swizzle_mx_scales(Ws) - _mark_static(Ws_fgm) # OpenAI triton_kernels dense mxfp4 matmul (matmul_ogs, no routing): the qkv linear # in the GPT-OSS MXFP4 format. Weight is a single (1, K, N) expert, swizzled once # at load (same as the fused arm); latency-only (its own weights). @@ -1509,7 +1459,6 @@ def _attn_arms(cfg, x, weights): tw_bf = torch.randn(1, K, N, device=DEVICE, dtype=torch.bfloat16) * 0.05 tw, tws = _tfmx.quantize_to_mxfp4(tw_bf, triton_kernels_hub) tw, tws = _tfmx.swizzle_mxfp4(tw, tws, triton_kernels_hub) - _mark_static(tw.storage.data, tws.storage.data) tk_pc = triton_kernels_hub.matmul_ogs.PrecisionConfig( weight_scale=tws, flex_ctx=triton_kernels_hub.matmul_ogs.FlexCtx( @@ -1540,7 +1489,6 @@ def _attn_arms(cfg, x, weights): rb = [ScalingType.BlockWise1x16, ScalingType.TensorWise] swz = [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE] mat_b = W.view(FP4).t() - _mark_static(mat_b, *sb, one) def torch_mm(a, a_scale=None): aq, a_s = fgm.nvfp4_act_quant(a) @@ -1553,7 +1501,6 @@ def torch_mm(a, a_scale=None): rb = ScalingType.BlockWise1x32 swz = SwizzleType.SWIZZLE_32_4_4 mat_b = W.t() - _mark_static(mat_b, sb) def torch_mm(a, a_scale=None): aq, a_s = fgm.mxfp8_act_quant(a) @@ -1569,7 +1516,6 @@ def torch_mm(a, a_scale=None): from torch.nn.functional import ScalingType mat_b = W.t() - _mark_static(mat_b) def torch_mm(a, a_scale): return torch.nn.functional.scaled_mm( @@ -1578,86 +1524,53 @@ def torch_mm(a, a_scale): Ws.reshape(()), ScalingType.TensorWise, output_dtype=torch.bfloat16) - # a STATIC row's calibrated activation scale replaces the inline quant: one value, since a - # dense module is one quantized module (the MoE's per-expert form has no dense analogue) - attn_act_scale = ( - (x.float().abs().amax() / 448.0).clamp(min=1e-12).reshape(1) - if cfg.get("static") else None - ) - if attn_act_scale is not None: - _mark_static(attn_act_scale) - # act is inline-quantized (As=None) unless the row is static; Ws is the weight scale (Bs); activation_format is the - # activation precision (None follows the weight format). - attn_arms = { - "finegrained-kernels": lambda: fgm.matmul_2d( - x, W, attn_act_scale, Ws_fgm, activation_format=_activation_format(cfg), output_dtype=torch.bfloat16, - b_global_scale=W_g), - } - if torch_mm is not None: - attn_arms["torch_mm"] = lambda: torch_mm(x, attn_act_scale) - if "finegrained-fp8" in cfg["baselines"]: - attn_arms["finegrained-fp8"] = lambda: upstream_fp8.matmul_2d(x, W, Ws_fp8, block, - torch.bfloat16) - if "deepgemm" in cfg["baselines"]: - attn_arms["deepgemm"] = lambda: deepgemm_fp8_fp4_linear( - x, W, Ws, block_size=dg_block, output_dtype=torch.bfloat16) - if "triton_kernels" in cfg["baselines"]: - attn_arms["triton_kernels"] = lambda: tk_ogs( - x, tw, None, None, precision_config=tk_pc) - # The dense GEMM is where transformers' NVFP4 path actually lives, so this is the - # apples-to-apples cell: same weight, its packer, its kernel. Packing is offline (the - # checkpoint conversion does it), so only `gemm` is timed. - if _nvfp4_gemm is not None and cfg["weights"] == "nvfp4": - # their packer takes a bf16 weight and produces its OWN layout, so feed it the - # dequantized tensor — `W` is already packed E2M1 (uint8, two values per byte), and - # casting those bytes to bf16 hands it a K/2-wide matrix ("expected 3072 input - # features, got 6144"). Same underlying weight, each stack quantizing it its own way, - # which is how every other baseline here is treated. - _w_bf16 = WEIGHTS["nvfp4"]["dequant"]( - W[None], Ws[None], (W_g if W_g is not None else torch.ones(1, device=W.device)) - )[0].to(torch.bfloat16) - _pw = _nvfp4_gemm.pack(_w_bf16, device=x.device) - _mark_static(_pw.qweight, _pw.sf, _pw.global_scale, _pw.sf_rowmajor) - attn_arms["nvfp4_gemm"] = lambda: _nvfp4_gemm.gemm(_pw, x) - return attn_arms - - -def bench_attn_row(row, pname, cfg, rows_out): - """One attn linear per model, in its deployment format (same weights across - impls; the finegrained-kernels arm's ``activation_format`` follows the model — GPT-OSS runs W4A4).""" - print(f"== [{row}] {pname}") - if MOCK: - _mock_rows(row, pname, ("finegrained-kernels",) + cfg["baselines"], rows_out) - return - torch._dynamo.reset() - N, K = cfg["N"], cfg["K"] - W_g = None - if cfg["weights"] == "fp8_128x128_ue8m0": - W, Ws = make_weights(N, K, DEVICE, [128, 128], - scale_dtype=torch.float8_e8m0fnu) - else: - # registry makers are expert-batched; build E=1 and index the slab off. NVFP4 returns a - # per-expert fp32 global as its 3rd value (None for single-level FP8/MX) — keep it so the - # two-level global multiply isn't dropped (it was, silently running NVFP4 at global=1). - W, Ws, W_g = WEIGHTS[cfg["weights"]]["make"](N, K, 1) - W, Ws = W[0], Ws[0] - W_g = W_g if W_g is None else W_g[0] - _mark_static(W, Ws) - weights = (W, Ws, W_g) - for regime, tokens in (("decode", DECODE_TOKENS), ("prefill", PREFILL_TOKENS)): print(f" -- {regime}") torch.manual_seed(0) x = torch.randn(tokens, K, device=DEVICE, dtype=torch.bfloat16) - banks, active_bytes = _weight_banks(weights) - inputs = [x] + [x.clone() for _ in banks[1:]] - before = ACCELERATOR.memory_allocated() - bank_arms = [_attn_arms(cfg, hidden, bank) for bank, hidden in zip(banks, inputs)] - prepared_bytes = max(0, ACCELERATOR.memory_allocated() - before) + # a STATIC row's calibrated activation scale replaces the inline quant: one value, since a + # dense module is one quantized module (the MoE's per-expert form has no dense analogue) + attn_act_scale = ( + (x.float().abs().amax() / 448.0).clamp(min=1e-12).reshape(1) + if cfg.get("static") else None + ) + if attn_act_scale is not None: + _mark_static(attn_act_scale) + # act is inline-quantized (As=None) unless the row is static; Ws is the weight scale (Bs); activation_format is the + # activation precision (None follows the weight format). + attn_arms = { + "finegrained-kernels": lambda: fgm.matmul_2d( + x, W, attn_act_scale, Ws_fgm, activation_format=_activation_format(cfg), output_dtype=torch.bfloat16, + b_global_scale=W_g), + } + if torch_mm is not None: + attn_arms["torch_mm"] = lambda: torch_mm(x, attn_act_scale) + if "finegrained-fp8" in cfg["baselines"]: + attn_arms["finegrained-fp8"] = lambda: upstream_fp8.matmul_2d(x, W, Ws_fp8, block, + torch.bfloat16) + if "deepgemm" in cfg["baselines"]: + attn_arms["deepgemm"] = lambda: deepgemm_fp8_fp4_linear( + x, W, Ws, block_size=dg_block, output_dtype=torch.bfloat16) + if "triton_kernels" in cfg["baselines"]: + attn_arms["triton_kernels"] = lambda: tk_ogs( + x, tw, None, None, precision_config=tk_pc) + # The dense GEMM is where transformers' NVFP4 path actually lives, so this is the + # apples-to-apples cell: same weight, its packer, its kernel. Packing is offline (the + # checkpoint conversion does it), so only `gemm` is timed. + if _nvfp4_gemm is not None and cfg["weights"] == "nvfp4": + # their packer takes a bf16 weight and produces its OWN layout, so feed it the + # dequantized tensor — `W` is already packed E2M1 (uint8, two values per byte), and + # casting those bytes to bf16 hands it a K/2-wide matrix ("expected 3072 input + # features, got 6144"). Same underlying weight, each stack quantizing it its own way, + # which is how every other baseline here is treated. + _w_bf16 = WEIGHTS["nvfp4"]["dequant"]( + W[None], Ws[None], (W_g if W_g is not None else torch.ones(1, device=W.device)) + )[0].to(torch.bfloat16) + _pw = _nvfp4_gemm.pack(_w_bf16, device=x.device) + attn_arms["nvfp4_gemm"] = lambda: _nvfp4_gemm.gemm(_pw, x) anchor_res, anchor_out = None, None - for name in bank_arms[0]: - runs = [arms[name] for arms in bank_arms] - res, out = bench_modes(runs, name, active_bytes, prepared_bytes) + for name, run in attn_arms.items(): + res, out = bench_modes(run, name) parity = None if anchor_res is None: anchor_res, anchor_out = res, out @@ -1668,12 +1581,11 @@ def bench_attn_row(row, pname, cfg, rows_out): print(f" {name:14s} parity-vs-finegrained-kernels {parity:.1e}" f" finegrained-kernels-speedup {sp}", flush=True) rows_out.append((row, pname, regime, name, res, parity)) - del runs, bank_arms, banks, inputs print() device_name = "MOCK (random values)" if MOCK else ACCELERATOR.get_device_name(0) -print(f"device: {device_name} torch {torch.__version__} cache={CACHE_POLICY}" +print(f"device: {device_name} torch {torch.__version__}" f"{' [SMOKE]' if SMOKE else ''}") print("finegrained-kernels = local build; baselines: finegrained-fp8 (upstream), DeepGEMM, " "vLLM fused MoE, FlashInfer's TRT-LLM routed MoE, transformers grouped_mm/batched_mm, " @@ -1689,8 +1601,7 @@ def wanted(*names): # owning one device via CUDA_VISIBLE_DEVICES). The coordinator spawns GPUS workers, each # writes a shard CSV, then the coordinator merges + plots. A single GPU (GPUS=1) runs inline. ── _CSV = os.path.join(_HERE, "bench_moe.csv") -_CSV_HEADER = ("category,problem,regime,impl,mode,latency_us,parity_vs_finegrained_kernels," - "cache_policy,bank_count,active_weight_bytes,prepared_bytes\n") +_CSV_HEADER = "category,problem,regime,impl,mode,latency_us,parity_vs_finegrained_kernels\n" def _write_rows_csv(path, rows_out): @@ -1701,9 +1612,7 @@ def _write_rows_csv(path, rows_out): v = res.get(mode) f.write(f'"{rr}","{p}",{reg},{i},{mode},' f'{"" if v is None else f"{v:.2f}"},' - f'{"" if par is None else par},{CACHE_POLICY},' - f'{res.get("bank_count", "")},{res.get("active_weight_bytes", "")},' - f'{res.get("prepared_bytes", "")}\n') + f'{"" if par is None else par}\n') def _load_rows_csv(path): @@ -1729,17 +1638,11 @@ def _allowed(cfg): allowed["linear quantized", pn] = _allowed(c) | {"torch_mm", "nvfp4_gemm"} acc = {} # (cat, problem, regime, impl) -> (res dict, parity) for r in csv.DictReader(open(path)): - if r.get("cache_policy") != CACHE_POLICY: - raise ValueError(f"{path}: missing/mismatched cache policy; rerun the benchmark " - "instead of mixing old warm-graph results") if r["impl"] not in allowed.get((r["category"], r["problem"]), {"finegrained-kernels"}): continue key = (r["category"], r["problem"], r["regime"], r["impl"]) res, par = acc.setdefault(key, ({}, None)) res[r["mode"]] = float(r["latency_us"]) if r["latency_us"] else None - for field in ("bank_count", "active_weight_bytes", "prepared_bytes"): - if r.get(field): - res[field] = int(r[field]) if r["parity_vs_finegrained_kernels"]: acc[key] = (res, float(r["parity_vs_finegrained_kernels"])) return [(cat, p, reg, impl, res, par) @@ -2033,10 +1936,9 @@ def _drawn(res): handles.append(mpatches.Patch(facecolor="none", edgecolor="black", hatch="..", label="eager (no cudagraph)")) ax.legend(handles=handles, loc="upper left", fontsize=8) -cache_note = "cold weights (>2xL2)" if DEVICE == "cuda" else "backend cache defaults" fig.suptitle(f"MoE bench — finegrained-kernels vs finegrained-fp8 + references " f"({device_name}, real model shapes; decode=cudagraph+compile, " - f"prefill=eager; {cache_note})", y=0.9995) + f"prefill=eager)", y=0.9995) fig.tight_layout(rect=(0, 0, 1, 0.99)) out_png = os.path.join(_HERE, f"bench_moe{suffix}.png") fig.savefig(out_png, dpi=120) diff --git a/finegrained-kernels/tests/test_autotuner.py b/finegrained-kernels/tests/test_autotuner.py index 53630fea6..78ae06ba6 100644 --- a/finegrained-kernels/tests/test_autotuner.py +++ b/finegrained-kernels/tests/test_autotuner.py @@ -19,7 +19,6 @@ import logging import subprocess import sys -import weakref import pytest import torch @@ -31,7 +30,6 @@ import finegrained_kernels # type: ignore import finegrained_kernels.matmul # type: ignore from finegrained_kernels.bayesian_autotuner import bayesian_autotune # type: ignore -from finegrained_kernels.bench import weight_input # type: ignore from finegrained_kernels.compat import is_sm10x # type: ignore from finegrained_kernels.pruners import block_within_dim_pruner, mx_config_pruner # type: ignore @@ -246,41 +244,6 @@ def _budget_kernel(X, Y, N, BLOCK: tl.constexpr, BAD: tl.constexpr): assert len(measured) == 3, _budget_kernel.configs_timings -@pytest.mark.kernels_ci -@pytest.mark.skipif(TEST_DEVICE != "cuda", reason="CUDA graph timing required") -@pytest.mark.parametrize("n_trials", [2, 3]) -def test_rotating_weights_are_reused_and_released(monkeypatch, n_trials): - """Cold graph tuning reuses identical banks across configs and frees them after the search.""" - weight = torch.randn(256, device=TEST_DEVICE) - output = torch.empty_like(weight) - monkeypatch.setattr("finegrained_kernels.bench.rotation_count", lambda active_bytes, device: 3) - pointers, references = {}, {} - - def observe(arguments): - value = arguments["B"] - pointers.setdefault(arguments["BLOCK"], set()).add(value.data_ptr()) - references[value.data_ptr()] = weakref.ref(value) - - @bayesian_autotune( - [triton.Config({"BLOCK": block}, pre_hook=observe) for block in (32, 64, 128)], - ["N"], - n_trials=n_trials, - benchmark_input=weight_input, - cache_results=False, - ) - @triton.jit - def copy_weight(B, C, N, BLOCK: tl.constexpr): - offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - tl.store(C + offsets, tl.load(B + offsets, offsets < N, 0), offsets < N) - - copy_weight[lambda meta: (triton.cdiv(256, meta["BLOCK"]),)](weight, output, 256) - torch.testing.assert_close(output, weight) - assert len(pointers) >= 2 - assert len({frozenset(values) for values in pointers.values()}) == 1 - assert all(len(values) == 3 for values in pointers.values()) - assert all(reference() is None for pointer, reference in references.items() if pointer != weight.data_ptr()) - - @pytest.mark.kernels_ci @pytest.mark.skipif(TEST_DEVICE is None, reason="accelerator (CUDA/XPU) required") def test_act_quant_arms_are_bit_equal(): diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/batched.py b/finegrained-kernels/torch-ext/finegrained_kernels/batched.py index 46cdee294..043e379fd 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/batched.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/batched.py @@ -22,7 +22,6 @@ from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait from .bayesian_autotuner import bayesian_autotune -from .bench import weight_input from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype, pdl_launch_kwargs from .descriptors import rebind_batched_mx_bs_descriptor @@ -105,7 +104,6 @@ def expert_setup( # two checkpoints at one shape with different blocks must not share a winner. ["N", "K", "S", "OUTPUT_FORMAT", "GATE", "BLOCK_N", "BLOCK_K"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ "early_config_prune": compose_pruners( @@ -250,7 +248,6 @@ def w8a8_block_dynamic_fp8_matmul_batched_kernel( # keyed like the block-dynamic sibling: requant narrows the legal tiles to the quant block ["N", "K", "S", "OUTPUT_FORMAT", "GATE", "BLOCK_N", "BLOCK_K"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ "early_config_prune": compose_pruners( @@ -385,7 +382,6 @@ def w8a8_block_static_fp8_matmul_batched_kernel( get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True), ["N", "K", "S"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is # row-masked only — veto non-dividing tiles on both. @@ -536,7 +532,6 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel( # incidentally today; this kernel's requant Cs is always row-major — no SWIZZLED_OUT axis). ["N", "K", "S", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BK-within-K + the sm_10x MMA-shape guards (swapped dot_scaled needs BN >= 128 for the # native scaled-MMA; smaller-BN swap configs never win and mislead the TPE). @@ -732,7 +727,6 @@ def mx_dynamic_matmul_batched_kernel( ), ["N", "K", "S", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ # No BK-within-K veto: this kernel's tail tile slides back in bounds (see the K-loop), @@ -873,7 +867,6 @@ def mx_weight_only_matmul_batched_kernel( # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide). ["N", "K", "S", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is # row-masked only — veto non-dividing tiles on both. diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py index f7958fa47..439c0e1dc 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/bayesian_autotuner.py @@ -37,7 +37,6 @@ import time from collections import defaultdict from contextlib import contextmanager -from functools import cached_property, partial from pathlib import Path import torch @@ -55,14 +54,10 @@ triton_key, ) -from .bench import do_bench, weight_banks - logger = logging.getLogger(__name__) -# Include timing and tuner implementations alongside Triton's kernel-source hash. -_TUNER_SOURCE_HASH = hashlib.sha256( - Path(__file__).read_bytes() + Path(__file__).with_name("bench.py").read_bytes() -).hexdigest() +# Include the tuner implementation alongside Triton's kernel-source hash. +_TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() # Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=): every @@ -101,15 +96,10 @@ def __init__( path_anchor_axes: tuple[str, ...] = (), coupled_axes: tuple[tuple[str, ...], ...] = (), finite_check_args: tuple[str, ...] = (), - benchmark_input=None, **kwargs, ): super().__init__(*args, **kwargs) self.fn_name = getattr(self.fn, "__name__", str(self.fn)) - # Large-model weights exceed L2; opt into repeated warm inputs for smaller workloads. - self.warm_cache = os.environ.get("FINEGRAINED_AUTOTUNE_WARM_CACHE", "0") == "1" - self.benchmark_input = benchmark_input - self._benchmark_banks = None # Bayesian trial budget — the per-decorator default, overridable via the # FINEGRAINED_AUTOTUNE_TRIALS env var (quick sweeps / exhaustive runs without touching # the decorators; set it >= grid size to fall back to stock exhaustive bench-all). @@ -157,18 +147,6 @@ def __init__( self._dtype_tag_cache: Dict[tuple, str] = {} self._dtype_str_cache: Dict[torch.dtype, str] = {} - @cached_property - def do_bench(self): - if self._do_bench is not None: - return self._do_bench - if ( - driver.active.get_current_target().backend != "cuda" - or self.user_defined_pre_hook or self.user_defined_post_hook - ): - return driver.active.get_benchmarker() - - return partial(do_bench, cudagraph=self.warm_cache) - def _dtype_str(self, dtype) -> str: s = self._dtype_str_cache.get(dtype) if s is None: @@ -210,15 +188,7 @@ def _bench(self, *args, config, **meta): self._failures.append((config, memo[sig] + " [memoized]")) return [float("inf")] * 3 try: - if ( - self.benchmark_input is not None and not self.warm_cache - and self._do_bench is None - and not self.user_defined_pre_hook and not self.user_defined_post_hook - and driver.active.get_current_target().backend == "cuda" - ): - timing = self._bench_rotating(*args, config=config, **meta) - else: - timing = super()._bench(*args, config=config, **meta) + timing = super()._bench(*args, config=config, **meta) bad = self._nonfinite_output(meta) if bad is not None: self._failures.append((config, f"non-finite output in {bad} (numerics veto)")) @@ -347,43 +317,6 @@ def _persist_failed_compile_memo(self): except Exception: pass # persistence is best-effort; the in-memory memo still holds - def _bench_rotating(self, *args, config, **meta): - if self._benchmark_banks is None: - arguments = {**dict(zip(self.arg_names, args)), **meta} - weight, active_bytes = self.benchmark_input(arguments) - mutated = self.reset_to_zero + self.restore_value - if any(arguments.get(name) is weight for name in mutated): - self._benchmark_banks = False - else: - self._benchmark_banks = weight_banks(arguments, weight, active_bytes) or False - if self._benchmark_banks is False: - logger.info("[autotune] %s: weight rotation unavailable; using cold event timing", self.fn_name) - if self._benchmark_banks is False: - return super()._bench(*args, config=config, **meta) - - config_args = config.all_kwargs() - conflicts = meta.keys() & config.kwargs.keys() - if conflicts: - raise ValueError(f"Conflicting meta-parameters: {', '.join(conflicts)}") - calls = [{**bank, **config_args} for bank in self._benchmark_banks] - - def run(): - for arguments in calls: - if config.pre_hook is not None: - config.pre_hook(arguments) - self.pre_hook(arguments) - try: - self.fn.run(**arguments) - except Exception as exc: - try: - self.post_hook(arguments, exception=exc) - finally: - raise - self.post_hook(arguments, exception=None) - - timing = do_bench(run, quantiles=(0.5, 0.2, 0.8), cudagraph=True) - return [value / len(calls) for value in timing] - def _report_bench_failures(self): """After every tune, report every UNIQUE failure — a failure is never silent: inf-scoring keeps the search alive, but a human must see what broke (e.g. a code @@ -444,10 +377,7 @@ def run(self, *args, **kwargs): if config_kwargs is None: config_kwargs = config._all_kwargs_memo = config.all_kwargs() return self.fn.run(*args, **kwargs, **config_kwargs) - try: - ret = super().run(*args, **kwargs) - finally: - self._benchmark_banks = None + ret = super().run(*args, **kwargs) self._report_bench_failures() return ret @@ -512,13 +442,10 @@ def benchmark(): full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()} self.pre_hook(full_nargs, reset_only=True) - try: - if self.cache_results: - self.check_disk_cache(key, pruned, benchmark) - else: - benchmark() - finally: - self._benchmark_banks = None + if self.cache_results: + self.check_disk_cache(key, pruned, benchmark) + else: + benchmark() config = self.cache[key] self.best_config = config @@ -854,11 +781,9 @@ def check_disk_cache(self, tuning_key, configs, bench_fn): # and FINEGRAINED_AUTOTUNE_TRIALS is not in Triton's invalidating-env list — # without n_trials a 3-trial exploratory winner gets replayed at 100 trials. str(self.keys), - str((getattr(self.benchmark_input, "__module__", None), - getattr(self.benchmark_input, "__qualname__", None))), f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}," f"runoff={self.runoff},basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters}," - f"coupled_axes={self.coupled_axes},warm_cache={self.warm_cache}", + f"coupled_axes={self.coupled_axes}", str(tuning_key), ] + [str(c) for c in configs] # str(Config) is pre_hook-free and process-stable cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest() @@ -1001,7 +926,6 @@ def bayesian_autotune( path_anchor_axes: tuple[str, ...] = (), coupled_axes: tuple[tuple[str, ...], ...] = (), finite_check_args: tuple[str, ...] = (), - benchmark_input=None, cache_results: bool = True, reset_to_zero=None, restore_value=None, @@ -1024,8 +948,6 @@ def bayesian_autotune( tuner stays independent of configuration details coupled_axes: groups of config fields treated as joint choices by TPE and refinement, including Config fields such as num_stages - benchmark_input: return a read-only weight tensor and its active read bytes - from the launch arguments for cold rotating graph timing max_failures: abort the tune once this many configs fail to compile/run before the measured budget is met (default n_trials; env override FINEGRAINED_AUTOTUNE_MAX_FAILURES). Without it a broadly broken @@ -1054,7 +976,6 @@ def decorator(fn): path_anchor_axes=path_anchor_axes, coupled_axes=coupled_axes, finite_check_args=finite_check_args, - benchmark_input=benchmark_input, cache_results=cache_results, **kwargs, ) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/bench.py b/finegrained-kernels/torch-ext/finegrained_kernels/bench.py deleted file mode 100644 index 9742c6ba1..000000000 --- a/finegrained-kernels/torch-ext/finegrained_kernels/bench.py +++ /dev/null @@ -1,133 +0,0 @@ -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Timing shared by autotuning and benchmarks.""" - -from dataclasses import replace - -import torch -from triton.tools.tensor_descriptor import TensorDescriptor -from triton.testing import do_bench as _do_bench -from triton.testing import do_bench_cudagraph - - -def do_bench(fn, quantiles=None, *, cudagraph=False, **kwargs): - """Time execution in milliseconds, excluding cache eviction. - - Event timing clears L2 before each measured interval. Graph timing has no - eviction: cold-weight callers must rotate a working set larger than L2 inside - ``fn``; otherwise it measures warm replay. - """ - if not cudagraph: - return _do_bench(fn, quantiles=quantiles, **kwargs) - # Triton's graph timer warms up on its own stream; finish input producers first. - torch.cuda.current_stream().synchronize() - return do_bench_cudagraph(fn, quantiles=quantiles, **kwargs) - - -def weight_input(nargs): - """Main GEMM weight and its logical read bytes, excluding inactive expert planes.""" - weight = nargs["B"] - active_bytes = weight.numel() * weight.element_size() - expert_ids = nargs.get("ExpertIds") - expert_start = nargs.get("ExpertStart") - if expert_ids is not None: - experts = weight.shape[0] - active = expert_ids[(expert_ids >= 0) & (expert_ids < experts)].unique().numel() - active_bytes = active_bytes * active // experts - elif expert_start is not None: - experts = weight.shape[0] - if expert_start.numel() < experts + 1: - return weight, 0 - active = (expert_start[1:experts + 1] > expert_start[:experts]).sum().item() - active_bytes = active_bytes * active // experts - return weight, active_bytes - - -def rotation_count(active_bytes, device): - """Cycle more than twice L2, or decline an impractically large captured graph.""" - if active_bytes <= 0: - return None - count = 2 * torch.cuda.get_device_properties(device).L2_cache_size // active_bytes + 1 - return count if count <= 1024 else None - - -def weight_banks(nargs, weight, active_bytes): - """Rotate readonly weight values; return None when cold event timing is safer. - - Activations, scales and outputs stay shared. Raw tensor aliases of the weight - are unsupported; host descriptors may alias views wholly inside its storage span. - """ - if not isinstance(weight, torch.Tensor) or not weight.is_cuda: - return None - if active_bytes > weight.numel() * weight.element_size(): - return None - # Broadcast/overlapping or holed layouts cannot count logical bytes as unique reads. - if torch._debug_has_internal_overlap(weight) != 0: - return None - names = [name for name, value in nargs.items() if value is weight] - if len(names) != 1: - return None - name = names[0] - storage = weight.untyped_storage()._cdata - descriptors = {} - span_bytes = weight.numel() * weight.element_size() - for key, value in nargs.items(): - if key != name and isinstance(value, torch.Tensor): - if value.untyped_storage()._cdata == storage: - return None - if isinstance(value, TensorDescriptor): - if not isinstance(value.base, torch.Tensor): - return None - if value.base.untyped_storage()._cdata != storage: - continue - base = value.base - offset = base.data_ptr() - weight.data_ptr() - extent = (1 + sum((size - 1) * stride for size, stride in zip(base.shape, base.stride()))) - if (offset < 0 or offset % 16 or offset % base.element_size() - or offset + extent * base.element_size() > span_bytes): - return None - descriptors[key] = (value, offset) - count = rotation_count(active_bytes, weight.device) - if count is None: - return None - free, _ = torch.cuda.mem_get_info(weight.device) - free += (torch.cuda.memory_reserved(weight.device) - - torch.cuda.memory_allocated(weight.device)) - if (count - 1) * span_bytes > free: - return None - banks = [nargs] - try: - with torch.no_grad(): - for _ in range(count - 1): - clone = torch.empty_strided(weight.shape, weight.stride(), - dtype=weight.dtype, device=weight.device) - clone.copy_(weight) - bank = dict(nargs, **{name: clone}) - rebound = {} - for key, (descriptor, offset) in descriptors.items(): - if id(descriptor) not in rebound: - base = descriptor.base - view = torch.empty(0, dtype=base.dtype, device=base.device).set_( - clone.untyped_storage(), offset // base.element_size(), - base.shape, base.stride(), - ) - rebound[id(descriptor)] = replace( - descriptor, base=view, block_shape=list(descriptor.block_shape), - ) - bank[key] = rebound[id(descriptor)] - banks.append(bank) - except torch.OutOfMemoryError: - return None - return banks diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py b/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py index 55cf77d6c..9423bdb28 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/grouped.py @@ -21,8 +21,7 @@ from triton.tools.tensor_descriptor import TensorDescriptor from .bayesian_autotuner import bayesian_autotune -from .bench import weight_input -from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, persistent_program_count, tl_dtype +from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, persistent_program_count, sm_count, tl_dtype from .descriptors import build_grouped_operand_descriptors, rebind_grouped_descriptors, rebind_grouped_mx_descriptors from .formats import check_activation_format, global_scale_stride, is_per_expert_global, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, routed_rows, tokens_per_expert_bucket, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, mx_act_quant_grouped, mx_act_quant_routed_rows, quantize_routed_rows_per_expert, static_expert_act_operands, swizzle_grouped_mx_scales, tensor_wide_act_operands @@ -64,7 +63,6 @@ # descriptor-box legality, so checkpoints with different blocks must not share a winner. ["N", "K", "tokens_per_expert_bit_length", "GATE", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # Pipeliner-race guard: per launch-BM, WS-only at BM >= 64 and non-WS below (see the pruner). prune_configs_by={ @@ -269,7 +267,6 @@ def w8a8_block_dynamic_fp8_matmul_grouped_kernel( ), ["N", "K", "tokens_per_expert_bit_length", "GATE", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # Pipeliner-race WS guard (per launch-BM: forces WS where the default pipeliner races at # BM>=64, keeps non-WS below) + descriptor-box limits. @@ -460,7 +457,6 @@ def w8a8_block_static_fp8_matmul_grouped_kernel( # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is # row-masked only — veto non-dividing tiles on both. WS is a pure perf axis here @@ -665,7 +661,6 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel( # splits them incidentally today; a descriptor Cs has no dtype token). ["N", "K", "tokens_per_expert_bit_length", "GATE", "SWIZZLED_SCALES", "OUTPUT_FORMAT", "SWIZZLED_OUT"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BK-within-K veto + the sm_10x dot_scaled shape/trap gates (this kernel had no # pruner while its BK span was {128,256} — the union span's BK=64 rows made the @@ -931,7 +926,6 @@ def mx_dynamic_matmul_grouped_kernel( ), ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, prune_configs_by={ "early_config_prune": compose_pruners( @@ -1101,7 +1095,6 @@ def mx_weight_only_matmul_grouped_kernel( # GATE keys the gate|up arm separately (its dot is 2*BN wide, a different tile optimum). ["N", "K", "tokens_per_expert_bit_length", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # BLOCK_SIZE_K is a tuned axis and the K-loop is maskless — veto non-dividing BKs. prune_configs_by={ diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py b/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py index f44d3b80c..2d624fa7d 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/matmul.py @@ -21,7 +21,6 @@ from triton.tools.tensor_descriptor import TensorDescriptor from .bayesian_autotuner import bayesian_autotune -from .bench import weight_input from .compat import add_op_namespace_prefix, FP8_DTYPE, is_sm10x, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype from .descriptors import maybe_descriptor, rebind_bd_descriptors, rebind_mx_descriptors, rebind_weight_only_descriptors from .formats import check_activation_format, global_scale_stride, normalize_global_scale, e2m1_as_uint8, is_mx, mx_scale_family, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_2d_operands, weight_format @@ -94,7 +93,6 @@ # at one shape with different blocks must not share a winner. ["N", "K", "m_bit_length", "GATE", "block_n", "block_k"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -285,7 +283,6 @@ def w8a8_block_dynamic_fp8_matmul_kernel( # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide, a distinct config space). ["N", "K", "m_bit_length", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -418,7 +415,6 @@ def w8a8_tensor_dynamic_fp8_matmul_kernel( # block_n/block_k (the launch-pinned quant block) gate descriptor-box legality and the K tile. ["N", "K", "m_bit_length", "GATE", "block_n", "block_k"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -592,7 +588,6 @@ def w8a8_block_static_fp8_matmul_kernel( # dtype token, so don't lean on the append). ["N", "K", "m_bit_length", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT", "SWIZZLED_OUT"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -792,7 +787,6 @@ def mx_dynamic_matmul_kernel( ), ["N", "K", "m_bit_length", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"), @@ -933,7 +927,6 @@ def mx_weight_only_matmul_2d_kernel( # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide). ["N", "K", "m_bit_length", "GATE"], n_trials=100, - benchmark_input=weight_input, path_anchor_axes=PATH_ANCHOR_AXES, # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe finite_check_args=("C", "Cs"),