From 61441461ecb778e7077dc6554a3f75469e0e1461 Mon Sep 17 00:00:00 2001 From: Changhao Wu Date: Thu, 13 Aug 2026 23:45:06 -0700 Subject: [PATCH 1/4] test(tilecpp): skip the two tilecpp cases ToT nvcc cannot build correctly --- tests/ops/test_bmm.py | 3 +++ tests/ops/test_rms_norm.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/tests/ops/test_bmm.py b/tests/ops/test_bmm.py index b296fec6..84093d01 100644 --- a/tests/ops/test_bmm.py +++ b/tests/ops/test_bmm.py @@ -71,6 +71,9 @@ def test_op( else: pytest.skip(f"Backend {backend} is not available") + if backend == "tilecpp" and static_persistent: + pytest.skip("tilecpp static_persistent is under investigation") + if backend == "cutile" and not static_persistent and (transpose_a or transpose_b): pytest.skip("CuTile non-persistent kernel doesn't support transpose") if backend == "cutile-rs" and not static_persistent and (transpose_a or transpose_b): diff --git a/tests/ops/test_rms_norm.py b/tests/ops/test_rms_norm.py index 35daf2c2..c129e384 100644 --- a/tests/ops/test_rms_norm.py +++ b/tests/ops/test_rms_norm.py @@ -62,6 +62,9 @@ def test_op(self, m, n, dtype, mode, backend, arch): if backend != "cutile" and mode == "multi_wave_cached": pytest.skip(f"multi_wave_cached mode is not implemented for backend {backend}") + if backend == "tilecpp" and mode == "static_persistent": + pytest.skip("tilecpp static_persistent is under investigation") + # skip static_persistent tests when n > 16384 to avoid excessive memory usage # Avoid tileiras hangs on RTX PRO 6000 which has 100 KB shared memory per SM # mode=None can also select static_persistent via heuristic when M > NUM_SMS * 2 From 292c84511392f6ffafccbe50dee362b851f60c0d Mon Sep 17 00:00:00 2001 From: Jinman Xie Date: Fri, 14 Aug 2026 00:52:42 -0700 Subject: [PATCH 2/4] refactor: unify autotune-disable env var to TILEGYM_DISABLE_AUTOTUNE --- .../suites/liger/cutile/fused_neighborhood_attention.py | 3 ++- src/tilegym/suites/liger/cutile/grpo_loss.py | 3 ++- tests/suites/liger/test_group_norm.py | 2 -- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/tilegym/suites/liger/cutile/fused_neighborhood_attention.py b/src/tilegym/suites/liger/cutile/fused_neighborhood_attention.py index 14ce29a6..a59c5b8b 100644 --- a/src/tilegym/suites/liger/cutile/fused_neighborhood_attention.py +++ b/src/tilegym/suites/liger/cutile/fused_neighborhood_attention.py @@ -14,6 +14,7 @@ from cuda.tile import RoundingMode as RMd from cuda.tile.tune import exhaustive_search +from tilegym.autotune import is_autotune_disabled from tilegym.backend import register_impl ConstInt = ct.Constant[int] @@ -604,7 +605,7 @@ def _fused_fwd_autotune( if cache_key not in _fwd_autotune_cache: configs = list(_fused_fwd_autotune_configs()) - if os.environ.get("DISABLE_AUTOTUNE", "0") == "1": + if is_autotune_disabled(): configs = configs[:1] def grid_fn(cfg): diff --git a/src/tilegym/suites/liger/cutile/grpo_loss.py b/src/tilegym/suites/liger/cutile/grpo_loss.py index 83426fa5..b548399b 100644 --- a/src/tilegym/suites/liger/cutile/grpo_loss.py +++ b/src/tilegym/suites/liger/cutile/grpo_loss.py @@ -32,6 +32,7 @@ LOG2E = 1.4426950408889634 +from tilegym.autotune import is_autotune_disabled from tilegym.backend import register_impl _LOSS_TYPE_GRPO = 0 @@ -361,7 +362,7 @@ def _grpo_loss_bwd_ct( def _tuned_fwd_kernel(stream, cache_key, grid, fwd_args): - if os.environ.get("DISABLE_AUTOTUNE") == "1": + if is_autotune_disabled(): return _grpo_loss_fwd_ct.replace_hints(occupancy=ByTarget(sm_100=_FWD_FALLBACK_OCC, default=_FWD_FALLBACK_OCC)) if cache_key not in _fwd_autotune_cache: result = exhaustive_search( diff --git a/tests/suites/liger/test_group_norm.py b/tests/suites/liger/test_group_norm.py index b86083c9..81c9dd38 100644 --- a/tests/suites/liger/test_group_norm.py +++ b/tests/suites/liger/test_group_norm.py @@ -37,7 +37,6 @@ def reference(X, num_channels, num_groups, W, B, eps=1e-5): @pytest.mark.parametrize("backend", _backends) def test_op_forward(self, batch_size, num_channels, hidden_size, num_groups, dtype, backend, monkeypatch): """Test forward output matches PyTorch F.group_norm reference.""" - monkeypatch.setenv("DISABLE_AUTOTUNE", "1") self.setUp() if tilegym.is_backend_available(backend): tilegym.set_backend(backend) @@ -72,7 +71,6 @@ def test_op_forward(self, batch_size, num_channels, hidden_size, num_groups, dty @pytest.mark.parametrize("backend", _backends) def test_op_backward(self, batch_size, num_channels, hidden_size, num_groups, dtype, backend, monkeypatch): """Test backward gradients (dX, dW, dB) match PyTorch reference.""" - monkeypatch.setenv("DISABLE_AUTOTUNE", "1") self.setUp() if tilegym.is_backend_available(backend): tilegym.set_backend(backend) From 3d39f669fd0c5bd53fadd849327ed36f264e6c29 Mon Sep 17 00:00:00 2001 From: Jinman Xie Date: Fri, 14 Aug 2026 00:55:29 -0700 Subject: [PATCH 3/4] add skip_perf_shape_on_oom decorator --- tests/common.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/common.py b/tests/common.py index 4d8ad5dc..562d8bf5 100644 --- a/tests/common.py +++ b/tests/common.py @@ -88,6 +88,38 @@ def get_tensor_alignment(tensor): return alignment +# Consumer-Blackwell arches (sm120/sm121) have less device memory than the +# data-center GPUs the liger ``test_perf`` shapes were sized for. +_MEMORY_CONSTRAINED_ARCHS = ("sm120", "sm121") + + +def skip_perf_shape_on_oom(test_fn): + r"""Convert a genuine OOM on memory-constrained arches into a skip. + + Wraps a liger ``test_perf`` method. If executing the perf shape raises a + ``torch.cuda.OutOfMemoryError`` and ``--arch`` is a memory-constrained + consumer-Blackwell arch (sm120/sm121), reclaim memory and ``pytest.skip``. + On every other arch (b200/sm100, h100/sm90, a100/sm80) the error is + re-raised so the perf-tracking platforms still fail loudly. Only + ``torch.cuda.OutOfMemoryError`` is intercepted -- correctness assertions and + all other exceptions propagate unchanged. + """ + + @wraps(test_fn) + def wrapper(self, *args, **kwargs): + try: + return test_fn(self, *args, **kwargs) + except torch.cuda.OutOfMemoryError: + gc.collect() + torch.cuda.empty_cache() + arch = self.request.config.getoption("--arch") + if arch in _MEMORY_CONSTRAINED_ARCHS: + pytest.skip(f"perf shape exceeds device memory on {arch}") + raise + + return wrapper + + class PyTestCase: r""" Base class for TileGym unit tests. From 8bee3fc0a70a116c49f11ed51685a9b11a0a6fd0 Mon Sep 17 00:00:00 2001 From: Jinman Xie Date: Mon, 17 Aug 2026 06:46:34 +0000 Subject: [PATCH 4/4] Added num_worker_warps to the cuTile persistent_layer_norm autotune configs --- src/tilegym/ops/cutile/layer_norm_legacy.py | 22 +++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/src/tilegym/ops/cutile/layer_norm_legacy.py b/src/tilegym/ops/cutile/layer_norm_legacy.py index 2389d85a..ff8a57ae 100644 --- a/src/tilegym/ops/cutile/layer_norm_legacy.py +++ b/src/tilegym/ops/cutile/layer_norm_legacy.py @@ -25,10 +25,15 @@ def _persistent_layer_norm_autotune_configs(): Generates configurations: - BLOCK_N: [1, 2, 4, 8, 16, 32] - number of rows per block + - num_worker_warps: [4, 8] - CUDA-core warp-group width (Triton ``num_warps`` + equivalent). Normalization-style kernels with large tiles are the + canonical case for tuning this hint, and nww=8 (256 threads) is a large + win on the bandwidth-bound small-D shapes. - num_ctas: [1] - single CTA for this kernel """ for block_n in [1, 2, 4, 8, 16, 32]: - yield SimpleNamespace(BLOCK_N=block_n, num_ctas=1) + for num_worker_warps in [4, 8]: + yield SimpleNamespace(BLOCK_N=block_n, num_ctas=1, num_worker_warps=num_worker_warps) def _get_default_persistent_layer_norm_configs(BLOCK_D=None): @@ -44,8 +49,8 @@ def _get_default_persistent_layer_norm_configs(BLOCK_D=None): block_n = min(8, p) else: block_n = 8 - return {"BLOCK_N": block_n, "num_ctas": 1} - return {"BLOCK_N": 8, "num_ctas": 1} + return {"BLOCK_N": block_n, "num_ctas": 1, "num_worker_warps": 8} + return {"BLOCK_N": 8, "num_ctas": 1, "num_worker_warps": 8} def _persistent_layer_norm_early_config_prune(configs, N, D, BLOCK_D): @@ -321,12 +326,14 @@ def grid_fn(cfg): grid_fn, _persistent_layer_norm_fwd_kernel, args_fn, - lambda cfg: {"num_ctas": cfg.num_ctas}, + lambda cfg: {"num_ctas": cfg.num_ctas, "num_worker_warps": cfg.num_worker_warps}, ) best_cfg = result.best.config _layer_norm_legacy_tune_cache[cache_key] = ( best_cfg, - _persistent_layer_norm_fwd_kernel.replace_hints(num_ctas=best_cfg.num_ctas), + _persistent_layer_norm_fwd_kernel.replace_hints( + num_ctas=best_cfg.num_ctas, num_worker_warps=best_cfg.num_worker_warps + ), ) best_cfg, tuned_kernel = _layer_norm_legacy_tune_cache[cache_key] ct.launch(stream, grid_fn(best_cfg), tuned_kernel, args_fn(best_cfg)) @@ -404,10 +411,13 @@ def _cutile_persistent_layer_norm_fwd( grid_size = min(NUM_SMS, num_row_blocks) grid = (grid_size, 1, 1) + default_kernel = _persistent_layer_norm_fwd_kernel.replace_hints( + num_ctas=configs["num_ctas"], num_worker_warps=configs["num_worker_warps"] + ) ct.launch( torch.cuda.current_stream(), grid, - _persistent_layer_norm_fwd_kernel, + default_kernel, ( x, y,