diff --git a/src/tilegym/ops/cutile/layer_norm_legacy.py b/src/tilegym/ops/cutile/layer_norm_legacy.py index 2389d85..ff8a57a 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, diff --git a/src/tilegym/suites/liger/cutile/fused_neighborhood_attention.py b/src/tilegym/suites/liger/cutile/fused_neighborhood_attention.py index 14ce29a..a59c5b8 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 83426fa..b548399 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/common.py b/tests/common.py index 4d8ad5d..562d8bf 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. diff --git a/tests/ops/test_bmm.py b/tests/ops/test_bmm.py index b296fec..84093d0 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 35daf2c..c129e38 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 diff --git a/tests/suites/liger/test_group_norm.py b/tests/suites/liger/test_group_norm.py index b86083c..81c9dd3 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)