From 7d033b1ed123f53802c4f83a44773b7664988ef0 Mon Sep 17 00:00:00 2001 From: Zhengyi Zhang Date: Wed, 23 Sep 2026 04:48:54 -0700 Subject: [PATCH 1/2] Unify cold-cache CUDA Graph measurement and autotuning Co-authored-by: Thomas Schmid --- src/tilegym/benchmark.py | 273 ++++++++ src/tilegym/ops/__init__.py | 2 - src/tilegym/ops/cutile/bmm.py | 4 +- .../ops/cutile/chunk_gated_delta_rule.py | 30 +- src/tilegym/ops/cutile/flash_decode.py | 5 +- src/tilegym/ops/cutile/group_gemm.py | 11 +- src/tilegym/ops/tilecpp/autotuner.py | 230 +++---- .../cutile/gemm/ragged_block_scaled_bmm.py | 2 +- .../flashinfer/cutile/gemm/ragged_bmm.py | 2 +- tests/autotune_cache.py | 44 ++ tests/common.py | 616 +++--------------- tests/config.py | 28 +- tests/conftest.py | 22 + tests/ops/activation/test_geglu.py | 2 +- tests/ops/test_attention.py | 4 +- tests/ops/test_attention_sink.py | 2 +- tests/ops/test_bmm.py | 47 +- tests/ops/test_chunk_gated_delta_rule.py | 2 +- tests/ops/test_dropout.py | 2 +- tests/ops/test_flash_decode.py | 2 +- tests/ops/test_gemma_attention.py | 2 +- tests/ops/test_gemma_attention_decode.py | 2 +- tests/ops/test_group_gemm.py | 2 +- tests/ops/test_layer_norm.py | 4 +- tests/ops/test_layer_norm_legacy.py | 4 +- tests/ops/test_linear_gluact_linear.py | 2 +- tests/ops/test_mamba2.py | 4 +- tests/ops/test_matmul.py | 6 +- tests/ops/test_mla.py | 9 +- tests/ops/test_mla_decoding.py | 2 +- tests/ops/test_mla_decoding_split_kv.py | 2 +- tests/ops/test_moe.py | 2 +- tests/ops/test_moe_actgrad_bwd.py | 2 - tests/ops/test_recurrent_gated_delta_rule.py | 2 +- tests/ops/test_rms_norm.py | 2 +- tests/ops/test_rope.py | 2 +- tests/ops/test_silu_and_mul.py | 2 +- tests/ops/test_splitk_reduce.py | 2 +- tests/ops/test_swiglu.py | 2 +- tests/ops/test_transpose.py | 2 +- 40 files changed, 630 insertions(+), 757 deletions(-) create mode 100644 src/tilegym/benchmark.py create mode 100644 tests/autotune_cache.py diff --git a/src/tilegym/benchmark.py b/src/tilegym/benchmark.py new file mode 100644 index 00000000..fea008df --- /dev/null +++ b/src/tilegym/benchmark.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: MIT + +"""Cold-cache CUDA Graph timing shared by benchmarks and autotuners.""" + +import math +import re +from contextlib import nullcontext + +import torch + +GRAPH_PROTOCOL = "cuda-graph-cold-span-v3" +FLUSH_BYTES = 256_000_000 +GRAPH_REPEATS = 8 + + +def iteration_counts(estimate_ms, warmup, rep, min_rep, max_rep): + if not math.isfinite(estimate_ms) or estimate_ms <= 0: + raise ValueError(f"Invalid GPU duration: {estimate_ms}") + n_warmup = max(1, int(warmup / estimate_ms)) + n_repeat = max(min_rep, int(rep / estimate_ms)) + if max_rep > 0: + n_repeat = min(n_repeat, max_rep) + n_warmup = min(n_warmup, max(1, int(max_rep * warmup / rep))) + return n_warmup, n_repeat + + +class _CapturedCall: + def __init__(self, fn, setup_fn, grad_to_none, fast_flush, *, cache=None, warmup=True, input_context=nullcontext): + self.graph = torch.cuda.CUDAGraph() + self.setup_graph = None + self.input_context = input_context + self.start = torch.cuda.Event(enable_timing=True, external=True) + self.end = torch.cuda.Event(enable_timing=True, external=True) + dtype = torch.int32 if fast_flush else torch.int8 + self.cache = ( + cache if cache is not None else torch.empty(FLUSH_BYTES // dtype.itemsize, dtype=dtype, device="cuda") + ) + caller = torch.cuda.current_stream() + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(caller) + try: + if warmup: + with torch.cuda.stream(capture_stream), self.input_context(): + if setup_fn is not None: + setup_fn() + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + fn() + capture_stream.synchronize() + pool = torch.cuda.graph_pool_handle() + with torch.cuda.stream(capture_stream), self.input_context(): + if setup_fn is not None: + self.setup_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.setup_graph, pool=pool, stream=capture_stream): + setup_fn() + with torch.cuda.graph(self.graph, pool=pool, stream=capture_stream): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + fn() + finally: + caller.wait_stream(capture_stream) + + def prepare(self): + if self.setup_graph is not None: + self.setup_graph.replay() + self.cache.zero_() + + def sample(self): + with self.input_context(): + self.prepare() + self.start.record() + self.graph.replay() + self.end.record() + self.end.synchronize() + return self.start.elapsed_time(self.end) + + def profile_samples(self, count, calls=None): + calls = [self] if calls is None else calls + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + ) as prof: + for index in range(count): + call = calls[index % len(calls)] + with call.input_context(): + call.prepare() + torch.cuda.synchronize() + with torch.profiler.record_function("tilegym.graph.sample"): + call.graph.replay() + torch.cuda.current_stream().synchronize() + # Flat timestamps avoid FunctionEvent construction and CPU call-tree processing. + scopes, device_events = [], [] + names = {} + for event in prof.profiler.kineto_results.events(): + if event.is_hidden_event(): + continue + device = event.device_type() + if device == torch.autograd.DeviceType.CPU and event.name() == "tilegym.graph.sample": + scopes.append((event.start_ns(), event.end_ns())) + elif device == torch.autograd.DeviceType.CUDA and not event.is_user_annotation(): + name = event.name() + if name not in names: + names[name] = torch._C._demangle(name) + device_events.append((event.start_ns(), event.end_ns(), names[name])) + scopes.sort(key=lambda event: event[0]) + device_events.sort(key=lambda event: event[0]) + if len(scopes) != count: + raise RuntimeError(f"Expected {count} profiler scopes, found {len(scopes)}") + samples = [] + index = 0 + for sample_index, (scope_start, scope_end) in enumerate(scopes): + while index < len(device_events) and device_events[index][1] <= scope_start: + index += 1 + selected = [] + while index < len(device_events) and device_events[index][0] < scope_end: + event = device_events[index] + if event[0] < scope_start or event[1] > scope_end: + raise RuntimeError("GPU activity crosses a synchronized sample boundary") + selected.append(event) + index += 1 + if not selected: + raise RuntimeError("CUDA Graph replay produced no profiled GPU activities") + samples.append( + { + "graph_index": sample_index % len(calls), + "span_ms": (max(event[1] for event in selected) - selected[0][0]) / 1_000_000, + "activities": [ + {"name": name, "duration_us": (end - start) / 1000} for start, end, name in selected + ], + } + ) + return samples + + +def _sample_kernel_records(samples): + records = {} + for sample in samples: + for activity in sample["activities"]: + name = activity["name"] + record = records.setdefault(name, {"name": name, "duration_us": 0.0, "count": 0}) + record["duration_us"] += activity["duration_us"] + record["count"] += 1 + return sorted( + [ + { + "name": record["name"], + "self_time_us": record["duration_us"] / len(samples), + "total_time_us": record["duration_us"] / len(samples), + "count": record["count"] / len(samples), + "sample_count": len(samples), + } + for record in records.values() + ], + key=lambda record: record["self_time_us"], + reverse=True, + ) + + +def benchmark_cuda_graph( + fn, + warmup=100.0, + rep=50.0, + min_rep=2, + initial_rep=5, + grad_to_none=None, + fast_flush=True, + *, + max_rep=1000, + setup_fn=None, + kernel_filter=None, + collect_kernel_times=False, + graph_repeats=GRAPH_REPEATS, + input_context=nullcontext, +): + """Measure cold-cache device spans with balanced samples across graphs. + + Compilation and tuning finish before capture. An optional setup graph runs + before the flush and timed graph, including when recreating backward inputs. + Capture errors propagate; this function never switches to eager execution. + CUPTI timestamps retain copies and gaps between GPU operations, excluding + timing-event overhead. A kernel filter explicitly sums matching activities. + Optional kernel metadata is averaged over the same timed replays. + Multiple independent captures include graph-instance and allocation effects + that repeated replay of one graph cannot estimate. Every graph contributes + the same number of samples; no graph or sample is selected for being faster. + An optional input context wraps every warmup, capture, and replay. Its entry + runs before the cache flush and its exit after timing, outside graph capture. + """ + if warmup < 0 or rep <= 0 or min_rep < 1 or initial_rep < 1 or max_rep < 0 or graph_repeats < 1: + raise ValueError("Invalid benchmark iteration budget") + if max_rep and max_rep < min_rep: + raise ValueError("max_rep must be zero or at least min_rep") + call = _CapturedCall(fn, setup_fn, grad_to_none, fast_flush, input_context=input_context) + estimate = sum(call.sample() for _ in range(initial_rep)) / initial_rep + n_warmup, n_repeat = iteration_counts(estimate, warmup, rep, min_rep, max_rep) + n_graphs = min(graph_repeats, n_repeat) + while True: + per_graph = math.ceil(n_repeat / n_graphs) + if max_rep: + per_graph = min(per_graph, max_rep // n_graphs) + if per_graph * n_graphs >= min_rep: + break + n_graphs -= 1 + n_repeat = per_graph * n_graphs + calls = [call] + for _ in range(n_graphs - 1): + calls.append( + _CapturedCall(fn, setup_fn, grad_to_none, fast_flush, cache=call.cache, input_context=input_context) + ) + for _ in range(math.ceil(n_warmup / n_graphs)): + for captured in calls: + with captured.input_context(): + captured.prepare() + captured.graph.replay() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + patterns = None + if kernel_filter is not None: + patterns = [re.compile(p) for p in ([kernel_filter] if isinstance(kernel_filter, str) else kernel_filter)] + profiled = call.profile_samples(n_repeat, calls=calls) + samples = [] + kernel_times = [] + for sample in profiled: + if patterns is None: + samples.append(sample["span_ms"]) + else: + matched = [k for k in sample["activities"] if any(p.search(k["name"]) for p in patterns)] + if not matched: + raise RuntimeError(f"No CUDA Graph kernels matched {kernel_filter!r}") + samples.append(sum(k["duration_us"] for k in matched) / 1000) + times = torch.tensor(samples, dtype=torch.float64) + mean = times.mean().item() + std = times.std(correction=0).item() + graph_means = torch.stack([times[index::n_graphs].mean() for index in range(n_graphs)]) + graph_std = graph_means.std(correction=0).item() + result = { + "mean": mean, + "std": std, + "rel_std": 100 * std / mean if mean else 0, + "median": times.median().item(), + "min": times.min().item(), + "max": times.max().item(), + "nrep": n_repeat, + "peak_mem_mb": torch.cuda.max_memory_allocated() // (1024 * 1024), + "measurement_protocol": GRAPH_PROTOCOL + ("-filtered" if patterns is not None else ""), + "cache_flush_bytes": FLUSH_BYTES, + "samples_ms": samples, + "graph_count": n_graphs, + "samples_per_graph": per_graph, + "graph_means_ms": graph_means.tolist(), + "graph_mean_std_ms": graph_std, + "graph_mean_rel_std": 100 * graph_std / mean if mean else 0, + } + if collect_kernel_times or patterns is not None: + kernel_times = _sample_kernel_records(profiled) + result["kernel_times"] = kernel_times + result["kernel_metadata_status"] = "recorded" if kernel_times else "unavailable" + result["kernel_metadata_source"] = "timed_samples" + return result + + +def autotune_cutile_cuda_graph(stream, grid, kernel, args): + import cuda.tile as ct + + # Upstream owns warmup; each callback must launch once on its prepared inputs. + with torch.cuda.stream(stream): + call = _CapturedCall( + lambda: ct.launch(torch.cuda.current_stream(), grid, kernel, args), None, None, True, warmup=False + ) + return call.profile_samples(1)[0]["span_ms"] * 1000 diff --git a/src/tilegym/ops/__init__.py b/src/tilegym/ops/__init__.py index 3a0b4b69..62586f40 100644 --- a/src/tilegym/ops/__init__.py +++ b/src/tilegym/ops/__init__.py @@ -66,8 +66,6 @@ from .attn_interface import mla_decoding_interface from .attn_interface import mla_interface from .moe_interface import fused_moe - -# Import all operation interfaces from the unified ops module from .ops import * __all__ = [ diff --git a/src/tilegym/ops/cutile/bmm.py b/src/tilegym/ops/cutile/bmm.py index bbce75e9..48924290 100644 --- a/src/tilegym/ops/cutile/bmm.py +++ b/src/tilegym/ops/cutile/bmm.py @@ -270,13 +270,13 @@ def grid_fn(cfg): num_ctas = getattr(cfg, "num_ctas", 1) base_programs = NUM_SMS // num_ctas - grid_size = min(base_programs, total_tiles) * occupancy + grid_size = min(base_programs * occupancy, total_tiles) return (grid_size,) # Call autotuner to find the best config and execute the kernel cache_key = (batch_size, M, N, K, transpose_a, transpose_b, a.dtype, str(a.device)) if cache_key not in _bmm_tune_cache: - with ct.compiler_timeout(15): + with ct.compiler_timeout(120): result = exhaustive_search( list(_bmm_autotune_configs()), stream, diff --git a/src/tilegym/ops/cutile/chunk_gated_delta_rule.py b/src/tilegym/ops/cutile/chunk_gated_delta_rule.py index 49e357e5..5cb96d60 100644 --- a/src/tilegym/ops/cutile/chunk_gated_delta_rule.py +++ b/src/tilegym/ops/cutile/chunk_gated_delta_rule.py @@ -68,18 +68,29 @@ def _ct_solve_tril_neumann_guarded(A, CS, n_steps): return result +@ct.kernel +def _chunk_cumsum_kernel(G, G_cum, NUM_HEADS: ConstInt, CHUNK_SIZE: ConstInt): + pid_bh = ct.bid(0) + pid_chunk = ct.bid(1) + b = pid_bh // NUM_HEADS + h = pid_bh % NUM_HEADS + g = ct.load(G, index=(b, pid_chunk, h), shape=(1, CHUNK_SIZE, 1), padding_mode=ct.PaddingMode.ZERO) + g = ct.astype(ct.reshape(g, (CHUNK_SIZE,)), ct.float32) + g_cum = ct.cumsum(g, axis=0) + ct.store(G_cum, index=(b, h, pid_chunk, 0), tile=ct.reshape(g_cum, (1, 1, 1, CHUNK_SIZE))) + + @ct.kernel(occupancy=2) def _intra_chunk_prepare_kernel( Q, K, V, Beta, - G, # raw inputs (B,T,H,D) / (B,T,H) Q_out, K_out, V_corr, K_cumdecay, - G_cum_out, # 5D outputs + G_cum, seq_len: int, K_dim: int, V_dim: int, @@ -123,16 +134,9 @@ def _intra_chunk_prepare_kernel( ct.float32, ) - g_raw = ct.astype( - ct.load( - G, - index=(b, pid_chunk, h), - shape=(1, CHUNK_SIZE, 1), - padding_mode=_ZERO, - ).reshape((CHUNK_SIZE,)), - ct.float32, + g_cum = ct.load(G_cum, index=(b, h, pid_chunk, 0), shape=(1, 1, 1, CHUNK_SIZE), padding_mode=_ZERO).reshape( + (CHUNK_SIZE,) ) - g_cum = ct.cumsum(g_raw, axis=0) offs_c = ct.arange(CHUNK_SIZE, dtype=ct.int32) offs_c_row = ct.expand_dims(offs_c, axis=1) @@ -186,7 +190,6 @@ def _intra_chunk_prepare_kernel( ) else: ct.store(K_cumdecay, index=(b, h, pid_chunk, 0, 0), tile=ct.reshape(kc_out, (1, 1, 1, CHUNK_SIZE, BLOCK_K))) - ct.store(G_cum_out, index=(b, h, pid_chunk, 0), tile=ct.reshape(g_cum, (1, 1, 1, CHUNK_SIZE))) q_tile = ct.astype( ct.load( @@ -394,6 +397,8 @@ def forward( output_buf = torch.empty(B, H, num_chunks, chunk_size, V, device=device, dtype=torch.float32) grid_intra = (B * H, num_chunks, 1) + # Materialize the scan before consumers transform it into MMA operand layouts. + ct.launch(torch.cuda.current_stream(), grid_intra, _chunk_cumsum_kernel, (g, g_cum, H, chunk_size)) ct.launch( torch.cuda.current_stream(), grid_intra, @@ -403,7 +408,6 @@ def forward( key, value, beta, - g, q_chunked, k_chunked, v_corrected, diff --git a/src/tilegym/ops/cutile/flash_decode.py b/src/tilegym/ops/cutile/flash_decode.py index 7e1679fc..b296938e 100644 --- a/src/tilegym/ops/cutile/flash_decode.py +++ b/src/tilegym/ops/cutile/flash_decode.py @@ -47,6 +47,7 @@ def attention_decode_kernel_grouped_impl( batch_id: int, head_id: int, tile_id: int, + ALLOW_TMA_STORE: ct.Constant[bool], ): """ cuTile device function for Grouped Query Attention decode with split-K parallelization. @@ -157,13 +158,12 @@ def attention_decode_kernel_grouped_impl( acc_reshaped = ct.reshape(acc, (1, 1, QUERY_GROUP_TILE_SIZE, 1, HEAD_DIM)) if NUM_Q_HEAD_PER_KV == QUERY_GROUP_TILE_SIZE: - # Use TMA store for optimal performance ct.store( Att_Out, index=(batch_id, head_id, 0, tile_id, 0), tile=acc_reshaped, order=(0, 1, 2, 3, 4), - allow_tma=True, + allow_tma=ALLOW_TMA_STORE, ) else: # Use scatter with boundary checking for non-matching tile sizes @@ -245,6 +245,7 @@ def _attention_decode_kernel_grouped( batch_id, head_id, tile_id, + True, ) diff --git a/src/tilegym/ops/cutile/group_gemm.py b/src/tilegym/ops/cutile/group_gemm.py index 2b4bf411..6e58c04a 100644 --- a/src/tilegym/ops/cutile/group_gemm.py +++ b/src/tilegym/ops/cutile/group_gemm.py @@ -37,9 +37,9 @@ def _group_gemm_autotune_configs(): @ct.kernel def _group_gemm_kernel( - As, # List of A matrices - Bs, # List of B matrices - Cs, # List of C matrices + As, + Bs, + Cs, TILE_M: ConstInt, TILE_N: ConstInt, TILE_K: ConstInt, @@ -51,7 +51,7 @@ def _group_gemm_kernel( group_size = len(As) zero_pad = ct.PaddingMode.ZERO - for g in range(group_size): + for g in ct.static_iter(range(group_size)): Ai = As[g] Bi = Bs[g] Ci = Cs[g] @@ -121,11 +121,12 @@ def _group_gemm_kernel( def _cutile_autotune_group_gemm(stream, group_A, group_B, group_C, transpose_b, device): """Autotune group GEMM kernel.""" + group_A, group_B, group_C = tuple(group_A), tuple(group_B), tuple(group_C) NUM_SMS = torch.cuda.get_device_properties(device).multi_processor_count group_shapes = tuple((tuple(A.shape), tuple(B.shape)) for A, B in zip(group_A, group_B)) cache_key = (group_shapes, transpose_b, group_A[0].dtype, str(group_A[0].device)) if cache_key not in _group_gemm_tune_cache: - with ct.compiler_timeout(5): + with ct.compiler_timeout(30): result = exhaustive_search( list(_group_gemm_autotune_configs()), stream, diff --git a/src/tilegym/ops/tilecpp/autotuner.py b/src/tilegym/ops/tilecpp/autotuner.py index b7027c7c..551fe61f 100644 --- a/src/tilegym/ops/tilecpp/autotuner.py +++ b/src/tilegym/ops/tilecpp/autotuner.py @@ -18,6 +18,7 @@ import logging import os import random +from contextlib import contextmanager from dataclasses import dataclass from typing import Any from typing import Callable @@ -26,6 +27,7 @@ import torch from tilegym.autotune import is_autotune_disabled +from tilegym.benchmark import benchmark_cuda_graph logger = logging.getLogger(__name__) @@ -124,53 +126,26 @@ def _time_ms( stream, warmup_ms: float = 25.0, rep_ms: float = 100.0, + setup_fn: Callable | None = None, ) -> float: - """Measure execution time in milliseconds using per-invocation CUDA events. + with torch.cuda.stream(stream): + return benchmark_cuda_graph(run_once, warmup=warmup_ms, rep=rep_ms, max_rep=96, setup_fn=setup_fn)["median"] - 1. Pilot run to estimate per-call cost. - 2. Derive warmup/repeat counts from time budgets. - 3. Per-invocation event pairs so each run is timed independently. - 4. Returns the **median** of a 10%-trimmed distribution for stability. - """ - stream.synchronize() - - # Pilot: estimate per-call cost - run_once() - stream.synchronize() - - pilot_start = torch.cuda.Event(enable_timing=True) - pilot_end = torch.cuda.Event(enable_timing=True) - pilot_start.record(stream) - for _ in range(5): - run_once() - pilot_end.record(stream) - pilot_end.synchronize() - estimate_ms = pilot_start.elapsed_time(pilot_end) / 5 - - n_warmup = max(1, int(warmup_ms / max(estimate_ms, 1e-3))) - n_repeat = max(10, int(rep_ms / max(estimate_ms, 1e-3))) - - # Warmup — stabilises GPU clocks, caches, and TLBs - for _ in range(n_warmup): - run_once() - stream.synchronize() - - # Benchmark with per-invocation events - starts = [torch.cuda.Event(enable_timing=True) for _ in range(n_repeat)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(n_repeat)] - for i in range(n_repeat): - starts[i].record(stream) - run_once() - ends[i].record(stream) - ends[-1].synchronize() - - times = sorted(s.elapsed_time(e) for s, e in zip(starts, ends)) - - # Trim fastest and slowest 10%, take median of the rest - lo = len(times) // 10 - hi = len(times) - lo - trimmed = times[lo:hi] if hi > lo else times - return trimmed[len(trimmed) // 2] + +@contextmanager +def _restore_inputs(tensors, stream): + with torch.cuda.stream(stream): + saved = [(tensor, tensor.clone()) for tensor in tensors] + + def restore(): + for tensor, snapshot in saved: + tensor.copy_(snapshot) + + try: + yield restore if saved else None + finally: + with torch.cuda.stream(stream): + restore() def _default_key( @@ -217,8 +192,9 @@ def __call__( named_args: dict[str, Any] = {}, *, max_iter: int = 60, - seed: int | None = None, + seed: int = 0, force_retune: bool = False, + restore_tensors: Sequence[torch.Tensor] = (), ) -> TunedResult: """ Run the autotuned kernel and return its result. @@ -242,13 +218,15 @@ def __call__( Maximum number of (valid) configurations to sample from the search space. seed: - Optional seed for the random number generator used when - sampling configurations. If ``None``, the global random number - generator state is used. + Seed for deterministic candidate ordering and sampling. force_retune: If ``True``, ignore any cached best config for this key and re-run the search. The new best config is then written back to the cache. + restore_tensors: + Inputs modified in place. Preserve their values during search; + reset before each candidate outside timing and the cache flush. + Cache hits launch directly without copying these inputs. Returns: TunedResult with the best configuration. @@ -261,84 +239,93 @@ def __call__( if verbose: logger.info(f"[TileCpp Autotuner] Cache hit for {key}: {best_cfg}") else: - if verbose: - logger.info(f"[TileCpp Autotuner] Starting autotuning for {key} with {len(self._search_space)} configs") - rng = random.Random(seed) - indices = rng.sample(range(len(self._search_space)), len(self._search_space)) + with _restore_inputs(restore_tensors, stream) as setup_fn: + if verbose: + logger.info( + f"[TileCpp Autotuner] Starting autotuning for {key} with {len(self._search_space)} configs" + ) + rng = random.Random(seed) + indices = rng.sample(range(len(self._search_space)), len(self._search_space)) + + # Phase 1: Pre-compile all configurations to warm up the compile cache + # This ensures JIT compilation overhead doesn't affect timing + if verbose: + logger.info(f"[TileCpp Autotuner] Pre-compiling up to {max_iter} configurations...") + valid_configs = [] + successes = 0 + for cfg_idx in indices: + if successes >= max_iter: + break + cfg = self._search_space[cfg_idx] + + # Apply filter predicate if defined + if not self._search_space.filter(named_args, cfg): + if verbose: + logger.debug(f"[TileCpp Autotuner] Config {cfg} filtered out by predicate function") + continue + + grid = grid_fn(named_args, cfg) + try: + # Run once to trigger JIT compilation + with torch.cuda.stream(stream): + if setup_fn is not None: + setup_fn() + launch_fn(cfg) + valid_configs.append((cfg_idx, cfg, grid)) + successes += 1 + except Exception as e: + if verbose: + logger.info(f"[TileCpp Autotuner] Config {cfg} failed during pre-compile: {e}") + continue - # Phase 1: Pre-compile all configurations to warm up the compile cache - # This ensures JIT compilation overhead doesn't affect timing - if verbose: - logger.info(f"[TileCpp Autotuner] Pre-compiling up to {max_iter} configurations...") - valid_configs = [] - successes = 0 - for cfg_idx in indices: - if successes >= max_iter: - break - cfg = self._search_space[cfg_idx] - - # Apply filter predicate if defined - if not self._search_space.filter(named_args, cfg): - if verbose: - logger.debug(f"[TileCpp Autotuner] Config {cfg} filtered out by predicate function") - continue - - grid = grid_fn(named_args, cfg) - try: - # Run once to trigger JIT compilation - launch_fn(cfg) - valid_configs.append((cfg_idx, cfg, grid)) - successes += 1 - except Exception as e: - if verbose: - logger.info(f"[TileCpp Autotuner] Config {cfg} failed during pre-compile: {e}") - continue - - if not valid_configs: - raise ValueError("No valid config found") - - # Synchronize to ensure all compilations are complete - stream.synchronize() - if verbose: - logger.info(f"[TileCpp Autotuner] Pre-compilation done. Timing {len(valid_configs)} valid configs...") + if not valid_configs: + raise ValueError("No valid config found") - # Phase 2: Time each pre-compiled configuration - best_time_ms, best_idx, best_grid = float("inf"), None, None + # Synchronize to ensure all compilations are complete + stream.synchronize() + if verbose: + logger.info( + f"[TileCpp Autotuner] Pre-compilation done. Timing {len(valid_configs)} valid configs..." + ) - for cfg_idx, cfg, grid in valid_configs: - try: + # Phase 2: Time each pre-compiled configuration + best_time_ms, best_idx, best_grid = float("inf"), None, None - def run_once(c=cfg): # Capture cfg in closure - launch_fn(c) + for cfg_idx, cfg, grid in valid_configs: + try: - time_ms = _time_ms(run_once, stream=stream) + def run_once(c=cfg): # Capture cfg in closure + launch_fn(c) - if time_ms < best_time_ms: - best_time_ms = time_ms - best_idx, best_grid = cfg_idx, grid - if verbose: - logger.info(f"[TileCpp Autotuner] New best: {cfg} -> {best_time_ms:.3f} ms") - else: - if verbose: - logger.info(f"[TileCpp Autotuner] Tried: {cfg} -> {time_ms:.3f} ms") + time_ms = _time_ms(run_once, stream=stream, setup_fn=setup_fn) - except Exception as e: - if verbose: - logger.info(f"[TileCpp Autotuner] Config {cfg} failed during timing: {e}") - continue + if time_ms < best_time_ms: + best_time_ms = time_ms + best_idx, best_grid = cfg_idx, grid + if verbose: + logger.info(f"[TileCpp Autotuner] New best: {cfg} -> {best_time_ms:.3f} ms") + else: + if verbose: + logger.info(f"[TileCpp Autotuner] Tried: {cfg} -> {time_ms:.3f} ms") - if best_idx is None: - raise ValueError("No valid config found after timing") + except Exception as e: + if verbose: + logger.info(f"[TileCpp Autotuner] Config {cfg} failed during timing: {e}") + continue - best_cfg = self._search_space[best_idx] - if verbose: - logger.info(f"[TileCpp Autotuner] Tuning complete. Best: {best_cfg} -> {best_time_ms:.3f} ms") - self._cache[key] = (best_idx, best_grid) + if best_idx is None: + raise ValueError("No valid config found after timing") + + best_cfg = self._search_space[best_idx] + if verbose: + logger.info(f"[TileCpp Autotuner] Tuning complete. Best: {best_cfg} -> {best_time_ms:.3f} ms") + self._cache[key] = (best_idx, best_grid) best_cfg = self._search_space[best_idx] # Launch with the best configuration - launch_fn(best_cfg) + with torch.cuda.stream(stream): + launch_fn(best_cfg) return TunedResult( best_cfg.kwargs, @@ -373,14 +360,7 @@ def wrapper(*args, **kwargs): def is_autotuning_enabled() -> bool: - """Report whether the search should run for this call. - - TILECPP_AUTOTUNE, when set, decides on its own so that this backend can be - pinned to its default configurations while the others keep searching. - Otherwise the project-wide TILEGYM_DISABLE_AUTOTUNE switch decides, which - leaves autotuning on by default. - """ - override = os.environ.get("TILECPP_AUTOTUNE") - if override is not None: - return override != "0" - return not is_autotune_disabled() + """Honor the global fixed-config policy before the backend-specific switch.""" + if is_autotune_disabled(): + return False + return os.environ.get("TILECPP_AUTOTUNE", "1") != "0" diff --git a/src/tilegym/suites/flashinfer/cutile/gemm/ragged_block_scaled_bmm.py b/src/tilegym/suites/flashinfer/cutile/gemm/ragged_block_scaled_bmm.py index 3fdd939f..a2879fd0 100644 --- a/src/tilegym/suites/flashinfer/cutile/gemm/ragged_block_scaled_bmm.py +++ b/src/tilegym/suites/flashinfer/cutile/gemm/ragged_block_scaled_bmm.py @@ -710,7 +710,7 @@ def ragged_block_scaled_bmm( # Materialize fallback max_m_device if the caller didn't pass one. The # kernel always reads its grid bound from a device tensor (defense-in-depth). if max_m_device is None: - max_m_device = torch.tensor([max_m], dtype=torch.int32, device=a.device) + max_m_device = torch.full((1,), max_m, dtype=torch.int32, device=a.device) has_a_scale = 1 if a_scale is not None else 0 if a_scale is None: diff --git a/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py b/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py index 5c7dcf81..a1ef5aba 100644 --- a/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py +++ b/src/tilegym/suites/flashinfer/cutile/gemm/ragged_bmm.py @@ -569,7 +569,7 @@ def ragged_bmm( # kernel always reads its grid bound from a device tensor (defense-in-depth), # so we keep the call sites uniform. if max_m_device is None: - max_m_device = torch.tensor([max_m], dtype=torch.int32, device=a.device) + max_m_device = torch.full((1,), max_m, dtype=torch.int32, device=a.device) # Check if autotune is enabled enable_autotune = is_autotune_enabled() diff --git a/tests/autotune_cache.py b/tests/autotune_cache.py new file mode 100644 index 00000000..6487e80e --- /dev/null +++ b/tests/autotune_cache.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: MIT + +import sys +from contextlib import contextmanager + + +def cutile_caches(): + """Include tuning choices stored in launch, hint and tile metadata caches.""" + for name, module in tuple(sys.modules.items()): + if module is None or not name.startswith("tilegym.") or "cutile" not in name.split("."): + continue + owners = [module] + owners.extend(value for value in vars(module).values() if isinstance(value, type) and value.__module__ == name) + for owner in owners: + for attr, value in vars(owner).items(): + if "cache" in attr.lower() and isinstance(value, dict): + yield owner, attr, value + + +def empty_cache(cache): + result = cache.copy() + result.clear() + return result + + +@contextmanager +def isolated_cutile_tune_caches(perf_caches): + original = {} + for owner, attr, cache in cutile_caches(): + key = (owner, attr) + original[key] = cache + if key not in perf_caches: + perf_caches[key] = empty_cache(cache) + setattr(owner, attr, perf_caches[key]) + try: + yield + finally: + for owner, attr, cache in cutile_caches(): + key = (owner, attr) + perf_caches[key] = cache + # Lazily imported operators have no pre-existing native scores. + setattr(owner, attr, original[key] if key in original else empty_cache(cache)) diff --git a/tests/common.py b/tests/common.py index d814c1d9..d3b78bc4 100644 --- a/tests/common.py +++ b/tests/common.py @@ -8,22 +8,18 @@ import gc import inspect import itertools -import json import multiprocessing import numbers import os import pathlib import random import re -import sys from functools import wraps import pytest - -# import pandas as pd import torch -from torch.profiler import ProfilerActivity -from torch.profiler import profile as torch_profile + +from tilegym.benchmark import benchmark_cuda_graph from .config import Config @@ -636,6 +632,8 @@ def find_modes_metrics(df): def load_previous(name, argument_names, expanded_argument_names): + import pandas as pd + load_dir = pathlib.Path(Config.load_dir) load_name = pathlib.Path(f"{name}.{Config.load}") load_path = load_dir / load_name @@ -956,7 +954,12 @@ def _compare_tensors_chunked( def any_output_requires_grad(fn): - out = fn() + caller = torch.cuda.current_stream() + stream = torch.cuda.Stream() + stream.wait_stream(caller) + with torch.cuda.stream(stream): + out = fn() + caller.wait_stream(stream) if not isinstance(out, tuple): out = (out,) @@ -973,10 +976,9 @@ def benchmark( initial_rep=Config.initial_rep, grad_to_none=None, fast_flush=True, - use_cudagraph=Config.cudagraph, - use_cupti=Config.cupti, kernel_filter=None, - retain_graph=True, + gradient=None, + setup_fn=None, ): r""" Benchmarks forward and backward pass of a given function. @@ -988,9 +990,9 @@ def benchmark( - ``'auto'`` - benchmarks forward pass and backward pass if any output requires gradient - Benchmarking of backward pass computes gradient with respect to the first - output argument. The function automatically generates input gradient w.r.t. - 1st output argument, gradient tensor is filled with ones. + By default backward differentiates the first output with an all-ones + gradient. Explicit gradients may be a tensor, a sequence matching the + outputs, or a callable receiving the outputs. Args: fn: function to be benchmarked @@ -1006,15 +1008,13 @@ def benchmark( every benchmark iteration fast_flush: if True: flush L2 cache with torch.int tensor, if False: flush L2 cache with torch.int8 tensor - kernel_filter: regex pattern to filter GPU kernels in CUPTI mode - (only used when use_cupti=True, ignored otherwise) - retain_graph: if True (default), the computation graph is retained - across backward iterations using ``retain_graph=True``. Set to - False for ops whose backward writes gradients in-place to saved - tensors (e.g. swiglu, geglu), which would corrupt the graph on - repeated calls. When False, the forward pass is re-run before - each backward iteration (not timed) so that only the backward - kernel time is measured. + kernel_filter: regex selecting GPU kernels in CUDA Graph replays + gradient: optional backward gradient tensor, sequence, or factory + setup_fn: optional input reset before each forward, outside timing and + before the cold-cache flush + Backward inputs are recreated in a separate setup graph before every + sample, outside the timed interval, on the same capture stream. Release + other live autograd graphs for these inputs before benchmarking. Returns: Dictionary with benchmark results for forward and backward pass @@ -1028,68 +1028,44 @@ def benchmark( ) ) - # For ops with in-place backward (e.g. swiglu, geglu): - result = common.benchmark( - lambda: tilegym.ops.swiglu(a, b), - retain_graph=False, - ) - """ - if kernel_filter is not None and not use_cupti: - raise ValueError("kernel_filter is only supported with use_cupti=True (CUPTI=1)") - if use_cupti: - # assert not use_cudagraph, "use_cudagraph and use_cupti cannot be True at the same time" - _benchmark_fn = functools.partial(benchmark_fn_cupti, kernel_filter=kernel_filter) - elif use_cudagraph: - _benchmark_fn = benchmark_fn_cudagraph - else: - _benchmark_fn = benchmark_fn + _benchmark_fn = functools.partial(benchmark_fn_cudagraph, kernel_filter=kernel_filter) res = {} if mode in ("auto", "forward"): - res["forward"] = _benchmark_fn(fn, warmup, rep, min_rep, initial_rep, grad_to_none, fast_flush) + res["forward"] = _benchmark_fn( + fn, warmup, rep, min_rep, initial_rep, grad_to_none, fast_flush, setup_fn=setup_fn + ) + if mode == "auto" and setup_fn is not None: + setup_fn() if mode == "backward" or (mode == "auto" and any_output_requires_grad(fn)): - if retain_graph: + state = {} + + def setup_backward(): + if setup_fn is not None: + setup_fn() out = fn() - if not isinstance(out, tuple): - out = (out,) - first_output = out[0] - assert first_output.requires_grad - dy = torch.ones_like(first_output) - bwd_fn = lambda: first_output.backward(dy, retain_graph=True) - res["backward"] = _benchmark_fn( - bwd_fn, - warmup, - rep, - min_rep, - initial_rep, - grad_to_none, - fast_flush, - ) - else: - # retain_graph=False: ops whose backward modifies saved tensors in-place - # (e.g. swiglu, geglu). Re-run forward in setup_fn (not timed) before - # each backward call so the graph is always fresh. - _state = {} - - def _setup_fn(): - out = fn() - if not isinstance(out, tuple): - out = (out,) - _state["first_output"] = out[0] - - def _bwd_fn(): - _state["first_output"].backward(torch.ones_like(_state["first_output"])) - - res["backward"] = benchmark_fn( - _bwd_fn, - warmup, - rep, - min_rep, - initial_rep, - grad_to_none, - fast_flush, - setup_fn=_setup_fn, - ) + if gradient is None: + if isinstance(out, (tuple, list)): + out = out[0] + grad = torch.ones_like(out) + else: + grad = gradient(out) if callable(gradient) else gradient + state["output"] = out + state["gradient"] = grad + + def backward(): + torch.autograd.backward(state["output"], state["gradient"]) + + res["backward"] = _benchmark_fn( + backward, + warmup, + rep, + min_rep, + initial_rep, + grad_to_none, + fast_flush, + setup_fn=setup_backward, + ) if grad_to_none is not None: for x in grad_to_none: x.grad = None @@ -1119,61 +1095,6 @@ def _pick_dominant_kernel(kernel_times, kernel_re_list=None): return candidates[0]["name"] if candidates else None -def _extract_kernel_times_from_profile(fn): - """ - Run fn() once under torch.profiler (CUDA only). Return the dominant kernel - name and a list of all CUDA kernel names with their times (for logging). - Returns (dominant_kernel_name, kernel_times_list) or (None, []) on failure. - kernel_times_list entries are dicts: name, self_time_us, total_time_us, count. - """ - if not torch.cuda.is_available(): - return None, [] - try: - with torch_profile(activities=[ProfilerActivity.CUDA]) as prof: - fn() - kernel_times = [] - for item in prof.key_averages(): - device_type_str = str(getattr(item, "device_type", "")) - if device_type_str != "DeviceType.CUDA": - continue - self_us = getattr(item, "self_device_time_total", 0.0) or 0.0 - total_us = getattr(item, "device_time_total", 0.0) or 0.0 - count = getattr(item, "count", 0) or 0 - kernel_times.append( - { - "name": item.key, - "self_time_us": self_us, - "total_time_us": total_us, - "count": int(count), - } - ) - # sort by self time descending so dominant kernel is first - kernel_times.sort(key=lambda x: x["self_time_us"], reverse=True) - return _pick_dominant_kernel(kernel_times), kernel_times - except Exception: - return None, [] - - -def _iteration_counts(estimate_ms, warmup, rep, min_rep, max_rep): - """Turn the millisecond budgets into iteration counts, bounded at both ends. - - ``warmup``/``rep`` are time budgets, so the iteration count scales with - ``1 / estimate_ms``. A kernel well below the ~50us launch-bound threshold - would otherwise inflate the loop into tens of thousands of profiled - launches, which costs wall-clock without adding signal -- the extra samples - only measure launch overhead. ``max_rep`` caps that; ``max_rep <= 0`` - restores the uncapped behaviour. The warmup cap keeps the original - ``warmup / rep`` ratio. - """ - n_warmup = max(1, int(warmup / estimate_ms)) - n_repeat = max(min_rep, int(rep / estimate_ms)) - if max_rep and max_rep > 0: - n_repeat = min(n_repeat, max_rep) - n_warmup = min(n_warmup, max(1, int(max_rep * warmup / rep))) - return n_warmup, n_repeat - - -# Adapted from https://github.com/openai/triton def benchmark_fn_cudagraph( fn, warmup=Config.warmup, @@ -1183,418 +1104,27 @@ def benchmark_fn_cudagraph( grad_to_none=None, fast_flush=True, max_rep=Config.max_rep, -): - n_retries = 10 - rep = rep / n_retries - # warmup - fn() - # extract kernel names and times from one profiled run (reuses warmup conceptually) - kernel_name, kernel_times = _extract_kernel_times_from_profile(fn) - if kernel_times and not Config.quiet: - print("Kernel times from profiled run (cudagraph):") - print(json.dumps(kernel_times, indent=2)) - # step 1 - we estimate the amount of time the kernel call takes - # NOTE: this is a single-run estimate, variance may be high. - if grad_to_none is not None: - for x in grad_to_none: - x.detach_() - x.requires_grad_(True) - x.grad = None - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - fn() - torch.cuda.synchronize() - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - g.replay() - end_event.record() - torch.cuda.synchronize() - estimate_ms = start_event.elapsed_time(end_event) - # The cap also bounds graph construction: `n_repeat` calls are unrolled into - # the graph below, so an uncapped count on a tiny kernel is paid twice. - _, n_repeat = _iteration_counts(estimate_ms, warmup, rep, min_rep, max_rep) - # step 2 - construct a cuda graph with `n_repeat` unrolled function calls to minimize - # host overhead - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - for i in range(n_repeat): - if grad_to_none is not None: - for x in grad_to_none: - x.grad = None - fn() - torch.cuda.synchronize() - # measure time and return - ret = [] - for i in range(n_retries): - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - g.replay() - end_event.record() - torch.cuda.synchronize() - ret += [start_event.elapsed_time(end_event) / n_repeat] - stats = torch.cuda.memory_stats() - peak_mem_mb = stats["allocated_bytes.all.peak"] // (1024 * 1024) - times = torch.tensor(ret) - res = { - "mean": times.mean().item(), - "std": times.std().item(), - "rel_std": (times.std() / times.mean()).item() * 100, - "median": times.median().item(), - "min": times.min().item(), - "max": times.max().item(), - "nrep": len(times), - "peak_mem_mb": peak_mem_mb, - } - if kernel_name is not None: - res["kernel_name"] = kernel_name - if kernel_times: - res["kernel_times"] = kernel_times - return res - - -# Adapted from https://github.com/openai/triton -def benchmark_fn( - fn, - warmup=Config.warmup, - rep=Config.rep, - min_rep=Config.min_rep, - initial_rep=Config.initial_rep, - grad_to_none=None, - fast_flush=True, setup_fn=None, - max_rep=Config.max_rep, -): - # setup_fn: optional callable run before each fn() call but outside the - # CUDA timing events, so only fn() is measured. Used to recreate a fresh - # computation graph when retain_graph=False (in-place backward ops). - # Estimate the runtime of the function - if setup_fn is not None: - setup_fn() - fn() - torch.cuda.synchronize() - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - for _ in range(initial_rep): - if setup_fn is not None: - setup_fn() - fn() - end_event.record() - torch.cuda.synchronize() - estimate_ms = start_event.elapsed_time(end_event) / initial_rep - # compute number of warmup and repeat - n_warmup, n_repeat = _iteration_counts(estimate_ms, warmup, rep, min_rep, max_rep) - # We maintain a buffer of 256 MB that we clear - # before each kernel call to make sure that the L2 - # doesn't contain any input data before the run - start_event = [torch.cuda.Event(enable_timing=True) for i in range(n_repeat)] - end_event = [torch.cuda.Event(enable_timing=True) for i in range(n_repeat)] - if fast_flush: - cache = torch.empty(int(256e6 // 4), dtype=torch.int, device="cuda") - else: - cache = torch.empty(int(256e6), dtype=torch.int8, device="cuda") - # Extract kernel names and times from one profiled warmup run - kernel_name, kernel_times = _extract_kernel_times_from_profile(fn) - if kernel_times and not Config.quiet: - print("Kernel times from profiled run:") - print(json.dumps(kernel_times, indent=2)) - # Warm-up - for _ in range(n_warmup): - if setup_fn is not None: - setup_fn() - fn() - # Benchmark - torch.cuda.empty_cache() - gc.collect() - torch._C._cuda_clearCublasWorkspaces() - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() - for i in range(n_repeat): - # we don't want `fn` to accumulate gradient values - # if it contains a backward pass. So we clear the - # provided gradients - if grad_to_none is not None: - for x in grad_to_none: - x.grad = None - # we clear the L2 cache before each run - cache.zero_() - torch.cuda._sleep(2000000) - # if setup_fn is provided, run it after cache flush but before timing - # so that only fn() is measured (used for ops with in-place backward) - if setup_fn is not None: - setup_fn() - # record time of `fn` - start_event[i].record() - fn() - end_event[i].record() - # Record clocks - torch.cuda.synchronize() - stats = torch.cuda.memory_stats() - peak_mem_mb = stats["allocated_bytes.all.peak"] // (1024 * 1024) - times = torch.tensor([s.elapsed_time(e) for s, e in zip(start_event, end_event)]) - q25 = times.quantile(0.25).item() - q75 = times.quantile(0.75).item() - res = { - "mean": times.mean().item(), - "std": times.std().item(), - "rel_std": (times.std() / times.mean()).item() * 100, - "median": times.median().item(), - "iqr": q75 - q25, - "q25": q25, - "q75": q75, - "min": times.min().item(), - "max": times.max().item(), - "nrep": len(times), - "peak_mem_mb": peak_mem_mb, - } - if kernel_name is not None: - res["kernel_name"] = kernel_name - if kernel_times: - res["kernel_times"] = kernel_times - return res - - -def _try_reset_kineto() -> None: - """Best-effort reset of the Kineto/CUPTI profiler state after an error. - - When an exception or external interrupt (e.g. ``pytest-timeout``) occurs - inside a ``torch.profiler.profile`` context, Kineto can be left in a - partially-initialised state. Subsequent tests that call - ``_prepare_profiler`` then fail with:: - - RuntimeError: Can't disable Kineto profiler when it's not running - - triggering a cascade of otherwise-unrelated failures across the entire - test session. - - This function attempts to complete the interrupted cleanup so the next - test starts with a clean Kineto slate. All errors are swallowed - intentionally: if the reset itself fails there is nothing more we can do, - and raising here would mask the original exception. - """ - for _reset in ( - # PyTorch ≥ 2.1 - lambda: torch._C._profiler._disable_profiler(), # type: ignore[attr-defined] - # Legacy fallback (PyTorch < 2.1) - lambda: torch._C._autograd._disable_profiler_legacy(), # type: ignore[attr-defined] - ): - try: - _reset() - return - except Exception: - pass - - -def benchmark_fn_cupti( - fn, - warmup=Config.warmup, - rep=Config.rep, - min_rep=Config.min_rep, - initial_rep=Config.initial_rep, - grad_to_none=None, - fast_flush=True, kernel_filter=None, - max_rep=Config.max_rep, ): - """ - Benchmark a function using CUPTI via ``torch.profiler``. - - Uses ``torch.profiler`` with ``ProfilerActivity.CUDA`` which collects - per-kernel GPU execution time through CUPTI callbacks — entirely in-process, - no ncu subprocess needed. - - Compared to CUDA Events: - - Measures pure kernel execution time (excludes launch overhead, CPU gaps) - - Each kernel is reported individually by name - - Args: - fn: function to benchmark (may launch one or more GPU kernels) - warmup: duration of warmup phase in milliseconds - rep: duration of measurement phase in milliseconds - min_rep: minimum number of measurement iterations - max_rep: maximum number of measurement iterations, so that a kernel far - below the launch-bound threshold does not turn the ``rep`` time - budget into tens of thousands of profiled launches; 0 disables - initial_rep: initial iterations to estimate runtime - grad_to_none: tensors whose ``.grad`` is set to None before each run - fast_flush: if True, flush L2 cache before each measurement - kernel_filter: filter for which GPU kernels to include in timing. - Can be a single regex string or a list of regex strings. - A kernel is included if it matches ANY of the patterns. - If None (default), all kernels are included. - Examples: - - ``"cutlass"`` — only CUTLASS kernels - - ``["softmax", "layer_norm"]`` — softmax or layer_norm kernels - - Returns: - dict compatible with ``benchmark_fn`` output: - ``{"mean", "std", "rel_std", "median", "min", "max", "nrep", "peak_mem_mb"}`` - """ - from torch.profiler import ProfilerActivity - from torch.profiler import profile - - if kernel_filter is None: - kernel_re_list = None - elif isinstance(kernel_filter, str): - kernel_re_list = [re.compile(kernel_filter)] - else: - kernel_re_list = [re.compile(p) for p in kernel_filter] - - # Step 1: estimate runtime to compute n_warmup and n_repeat - fn() - torch.cuda.synchronize() - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - for _ in range(initial_rep): - fn() - end_event.record() - torch.cuda.synchronize() - estimate_ms = start_event.elapsed_time(end_event) / initial_rep - - n_warmup, n_repeat = _iteration_counts(estimate_ms, warmup, rep, min_rep, max_rep) - - # Step 2: warmup - for _ in range(n_warmup): - fn() - torch.cuda.synchronize() - - # Step 3: L2 cache flush setup (same as benchmark_fn) - if fast_flush: - cache = torch.empty(int(256e6 // 4), dtype=torch.int, device="cuda") - else: - cache = torch.empty(int(256e6), dtype=torch.int8, device="cuda") - - # Step 4: measure with torch.profiler (CUPTI) - torch.cuda.empty_cache() - gc.collect() - torch._C._cuda_clearCublasWorkspaces() - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() - - run_times_us = [] - for i in range(n_repeat): - if grad_to_none is not None: - for x in grad_to_none: - x.grad = None - - # Flush L2 before enabling CUPTI, then wait for the fill kernel to - # finish so it cannot be included in the profiled kernel times. - cache.zero_() - torch.cuda.synchronize() - - try: - with profile(activities=[ProfilerActivity.CUDA]) as prof: - fn() - torch.cuda.synchronize() - except Exception: - # The profiler may have been left in a broken Kineto state, e.g. - # after a pytest-timeout interrupts _disable_profiler mid-flight, - # or after a GPU kernel hang. Attempt cleanup here so the *next* - # test case can still use CUPTI instead of cascading into dozens - # of unrelated "Can't disable Kineto profiler" failures. - _try_reset_kineto() - raise - - # Sum all kernel times for this run (fn may launch multiple kernels) - total_us = sum( - evt.self_device_time_total - for evt in prof.key_averages() - if evt.self_device_time_total > 0 - and (kernel_re_list is None or any(r.search(evt.key) for r in kernel_re_list)) - ) - - # DEBUG: print included events on the middle iteration (skip cold-start noise) - if i == n_repeat // 2 and os.environ.get("DUMP_CUPTI_EVENTS", "0") == "1": - print(f"=== CUPTI events (total={total_us:.1f} us) ===", file=sys.stderr) - for evt in prof.key_averages(): - if evt.self_device_time_total > 0 and ( - kernel_re_list is None or any(r.search(evt.key) for r in kernel_re_list) - ): - per_call = evt.self_device_time_total / evt.count - print( - f" {evt.key[:100]} " - f"total={evt.self_device_time_total:.1f} us " - f"per_call={per_call:.2f} us rep_num={evt.count}", - file=sys.stderr, - ) - print("=== end ===", file=sys.stderr) - - # On the first iteration, diagnose why total_us might be 0. - # Three distinct failure modes: - # 1. CUPTI/libcupti not available — torch.profiler silently returns 0 - # 2. fn() launched no GPU kernels — nothing to measure - # 3. kernel_filter regex matched none of the actual kernels - if i == 0 and total_us == 0: - all_kernel_us = sum( - evt.self_device_time_total for evt in prof.key_averages() if evt.self_device_time_total > 0 - ) - if all_kernel_us == 0: - # Could be (1) or (2) — check if ANY cuda events were recorded - has_cuda_events = any( - "cuda" in evt.key.lower() or evt.device_time_total > 0 for evt in prof.key_averages() - ) - if not has_cuda_events: - raise RuntimeError( - "CUPTI returned 0 device time — libcupti may not be available. " - "Check that nvidia-cuda-cupti is installed, or unset the CUPTI env var." - ) - else: - raise RuntimeError( - "fn() did not launch any GPU kernels. " - "Ensure the benchmarked function runs at least one CUDA kernel." - ) - if kernel_re_list is not None: - matched_names = [evt.key for evt in prof.key_averages() if evt.self_device_time_total > 0] - raise RuntimeError( - f"kernel_filter matched no GPU kernels. Filter: {kernel_filter}, actual kernels: {matched_names}" - ) - - run_times_us.append(total_us) - - times = torch.tensor(run_times_us, dtype=torch.float64) / 1000.0 # us -> ms - - stats = torch.cuda.memory_stats() - peak_mem_mb = stats["allocated_bytes.all.peak"] // (1024 * 1024) - - # Extract kernel names and times from the last profiler run. - # ``prof`` is still in scope here (the last loop iteration's profiler), so we - # can call key_averages() without launching an extra profiled run. - # All kernels with self_device_time_total > 0 are captured regardless of - # kernel_filter, giving a complete picture of every GPU kernel that ran. - # kernel_name however must describe what was timed: it is picked from the - # kernel_filter matches (when set), skipping helper kernels. - cupti_kernel_times = [] - for _item in prof.key_averages(): - if _item.self_device_time_total > 0: - cupti_kernel_times.append( - { - "name": _item.key, - "self_time_us": _item.self_device_time_total, - "total_time_us": _item.device_time_total, - "count": int(_item.count), - } - ) - cupti_kernel_times.sort(key=lambda x: x["self_time_us"], reverse=True) - kernel_name = _pick_dominant_kernel(cupti_kernel_times, kernel_re_list) - - res = { - "mean": times.mean().item(), - "std": times.std().item(), - "rel_std": (times.std() / times.mean()).item() * 100 if times.mean().item() > 0 else 0, - "median": times.median().item(), - "min": times.min().item(), - "max": times.max().item(), - "nrep": len(times), - "peak_mem_mb": peak_mem_mb, - } - if kernel_name is not None: - res["kernel_name"] = kernel_name - if cupti_kernel_times: - res["kernel_times"] = cupti_kernel_times - return res + result = benchmark_cuda_graph( + fn, + warmup, + rep, + min_rep, + initial_rep, + grad_to_none, + fast_flush, + max_rep=max_rep, + setup_fn=setup_fn, + kernel_filter=kernel_filter, + collect_kernel_times=True, + ) + patterns = None + if kernel_filter is not None: + patterns = [re.compile(p) for p in ([kernel_filter] if isinstance(kernel_filter, str) else kernel_filter)] + result["kernel_name"] = _pick_dominant_kernel(result["kernel_times"], patterns) + return result def benchmark_framework(framework_name, framework_fn, **benchmark_kwargs): diff --git a/tests/config.py b/tests/config.py index f13aefb4..f502c8be 100644 --- a/tests/config.py +++ b/tests/config.py @@ -56,6 +56,9 @@ def __new__(cls): @classmethod def parse(cls): + for removed in ("CUPTI", "CUDAGRAPH"): + if removed in os.environ: + raise ValueError(f"Remove {removed}: performance tests always use CUDA Graphs") parser = argparse.ArgumentParser( prog="TileGym Tests", description="TileGym Tests", @@ -171,10 +174,7 @@ def parse(cls): type=int, help=( "set MAX_REP to specify the maximum number of measured iterations for " - "performance tests. REP is a time budget, so a kernel far below the " - "~50us launch-bound threshold would otherwise inflate the loop into tens " - "of thousands of profiled launches without adding signal. Set to 0 to " - "disable the cap" + "performance tests. Set to 0 to disable the cap" ), ) parser.add_argument( @@ -215,24 +215,6 @@ def parse(cls): type=bool, help=("Record to csv file"), ) - parser.add_argument( - "--cudagraph", - envvar="CUDAGRAPH", - action=FromEnvironment, - required=False, - default=False, - type=bool, - help=("Use cudagraph"), - ) - parser.add_argument( - "--cupti", - envvar="CUPTI", - action=FromEnvironment, - required=False, - default=True, - type=bool, - help=("Use CUPTI (torch.profiler) for kernel profiling instead of CUDA Events"), - ) parser.add_argument( "--file", envvar="FILE", @@ -240,7 +222,7 @@ def parse(cls): required=False, default="out", type=str, - help=("When use cudagraph & csv, can specify file path"), + help=("Output file path when recording CSV"), ) parser.add_argument( "--config", diff --git a/tests/conftest.py b/tests/conftest.py index 5a880149..1cf41767 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,11 @@ import pytest +from tests.autotune_cache import isolated_cutile_tune_caches + +_graph_cutile_autotune_caches = {} + + try: import torch except ImportError: @@ -107,6 +112,23 @@ def framework(request): return request.config.getoption("--framework") +@pytest.fixture(autouse=True) +def _graph_cutile_autotune_in_perf_tests(request, monkeypatch): + if "test_perf" not in request.node.name: + yield + return + try: + from cuda.tile.tune import _tune + except ImportError: + yield + return + from tilegym.benchmark import autotune_cutile_cuda_graph + + with isolated_cutile_tune_caches(_graph_cutile_autotune_caches): + monkeypatch.setattr(_tune, "_benchmark", autotune_cutile_cuda_graph) + yield + + def _has_object_repr(val): """Helper function to recursively check if any value in nested structure has object representation""" if isinstance(val, (int, float, str, bool, type(None))): diff --git a/tests/ops/activation/test_geglu.py b/tests/ops/activation/test_geglu.py index be6b514e..6494b38d 100644 --- a/tests/ops/activation/test_geglu.py +++ b/tests/ops/activation/test_geglu.py @@ -167,5 +167,5 @@ def test_perf(self, x_shape, dim, dtype, approximate, framework, record_property atol=1e-2, ) - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=False) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 36260af4..d7d5ab79 100644 --- a/tests/ops/test_attention.py +++ b/tests/ops/test_attention.py @@ -195,7 +195,7 @@ def _run_perf(self, batch, heads, seq_len, head_dim, dtype, is_causal, backend, rtol=rtol, check_stride=False, ) - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=True) + result = common.benchmark_framework(backend, backend_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM @@ -244,7 +244,7 @@ def _run_perf_llm(self, model, batch_size, num_heads, seq_len, head_dim, dtype, atol=atol, check_stride=False, ) - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=True) + result = common.benchmark_framework(backend, backend_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_attention_sink.py b/tests/ops/test_attention_sink.py index 7606a476..aa1b1a28 100644 --- a/tests/ops/test_attention_sink.py +++ b/tests/ops/test_attention_sink.py @@ -217,7 +217,7 @@ def test_perf( check_stride=False, ) - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=True) + result = common.benchmark_framework(backend, backend_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_bmm.py b/tests/ops/test_bmm.py index 84093d01..613eb379 100644 --- a/tests/ops/test_bmm.py +++ b/tests/ops/test_bmm.py @@ -2,7 +2,9 @@ # # SPDX-License-Identifier: MIT +import importlib import itertools +from types import SimpleNamespace import pytest import torch @@ -22,6 +24,49 @@ def reference(a, b, transpose_a=False, transpose_b=False): b = torch.transpose(b, 1, 2) return torch.bmm(a, b) + @pytest.mark.parametrize("backend", ["cutile"]) + @pytest.mark.parametrize("m,n,k", [(128, 256, 511), (129, 257, 65), (1024, 512, 1023)]) + @pytest.mark.parametrize("transpose_a,transpose_b", [(False, True), (True, False)]) + def test_op_static_persistent_cudagraph(self, backend, m, n, k, transpose_a, transpose_b, monkeypatch): + try: + tilegym.set_backend(backend) + except Exception as e: + pytest.skip(f"Backend is not supported: {e}") + self.setUp() + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("This multi-CTA regression targets SM100-family GPUs") + + module = importlib.import_module("tilegym.ops.cutile.bmm") + config = SimpleNamespace(TILE_M=128, TILE_N=256, TILE_K=64, GROUP_SIZE_M=8, LATENCY=3, occupancy=2, num_ctas=2) + a = torch.rand((4, k, m) if transpose_a else (4, m, k), dtype=torch.float16, device="cuda") + b = torch.rand((4, n, k) if transpose_b else (4, k, n), dtype=torch.float16, device="cuda") + key = (4, m, n, k, transpose_a, transpose_b, a.dtype, str(a.device)) + kernel = module._static_persistent_bmm_kernel.replace_hints(num_ctas=2, occupancy=2) + monkeypatch.setattr(module, "_bmm_tune_cache", {key: (config, kernel)}) + + def run(): + return tilegym.ops.bmm(a, b, transpose_a=transpose_a, transpose_b=transpose_b, static_persistent=True) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = run() + torch.cuda.current_stream().wait_stream(stream) + + for _ in range(3): + a.uniform_() + b.uniform_() + expected = self.reference(a, b, transpose_a=transpose_a, transpose_b=transpose_b) + graph.replay() + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-8) + _backends = ["cutile"] if is_backend_available("tilecpp"): _backends = _backends + ["tilecpp"] @@ -157,7 +202,7 @@ def test_perf(self, q, m, n, k, transpose_a, transpose_b, dtype, framework, reco atol=1e-8, ) - res = common.benchmark_framework(framework, framework_fn, use_cudagraph=False) + res = common.benchmark_framework(framework, framework_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_chunk_gated_delta_rule.py b/tests/ops/test_chunk_gated_delta_rule.py index 659c23a0..2d4d0fe3 100644 --- a/tests/ops/test_chunk_gated_delta_rule.py +++ b/tests/ops/test_chunk_gated_delta_rule.py @@ -278,7 +278,7 @@ def test_perf(self, T, framework, record_property, arch): else: pytest.skip(f"Framework {framework} is not available") - result = common.benchmark_framework(framework, framework_fn, use_cupti=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) del q, k, v, g, beta diff --git a/tests/ops/test_dropout.py b/tests/ops/test_dropout.py index 216b95c3..a56e9f43 100644 --- a/tests/ops/test_dropout.py +++ b/tests/ops/test_dropout.py @@ -121,7 +121,7 @@ def test_perf(self, m, p, inplace, dtype, framework, record_property): else: pytest.skip(f"Framework {framework} is not available") - res = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + res = common.benchmark_framework(framework, framework_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_flash_decode.py b/tests/ops/test_flash_decode.py index e0cf0414..445badc4 100644 --- a/tests/ops/test_flash_decode.py +++ b/tests/ops/test_flash_decode.py @@ -183,7 +183,7 @@ def test_perf( rtol=rtol, check_stride=False, ) - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_gemma_attention.py b/tests/ops/test_gemma_attention.py index 1d8cc509..15d92fa7 100644 --- a/tests/ops/test_gemma_attention.py +++ b/tests/ops/test_gemma_attention.py @@ -354,7 +354,7 @@ def test_perf( ) # Benchmark - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Cleanup diff --git a/tests/ops/test_gemma_attention_decode.py b/tests/ops/test_gemma_attention_decode.py index a78eb205..924d05cf 100644 --- a/tests/ops/test_gemma_attention_decode.py +++ b/tests/ops/test_gemma_attention_decode.py @@ -297,7 +297,7 @@ def test_perf( ) # Benchmark - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=True) + result = common.benchmark_framework(backend, backend_fn) record_property("benchmark", result) # Cleanup diff --git a/tests/ops/test_group_gemm.py b/tests/ops/test_group_gemm.py index 2dc5c71b..90781c13 100644 --- a/tests/ops/test_group_gemm.py +++ b/tests/ops/test_group_gemm.py @@ -163,7 +163,7 @@ def test_perf( res = backend_fn() return - res = common.benchmark_framework(backend, backend_fn, use_cudagraph=False) + res = common.benchmark_framework(backend, backend_fn) record_property("benchmark", res) if dtype == torch.float8_e5m2: atol = 1 diff --git a/tests/ops/test_layer_norm.py b/tests/ops/test_layer_norm.py index 10f182dd..bfd7e7d0 100644 --- a/tests/ops/test_layer_norm.py +++ b/tests/ops/test_layer_norm.py @@ -101,7 +101,7 @@ def test_perf(self, m, n, k, dtype, framework, weight_shift, record_property): else: pytest.skip(f"Framework {framework} not supported") with torch.no_grad(): - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM @@ -204,7 +204,7 @@ def test_perf(self, m, n, dtype, framework, weight_shift, record_property): else: pytest.skip(f"Framework {framework} not supported") with torch.no_grad(): - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_layer_norm_legacy.py b/tests/ops/test_layer_norm_legacy.py index 6c4b55dc..2a56a2a2 100644 --- a/tests/ops/test_layer_norm_legacy.py +++ b/tests/ops/test_layer_norm_legacy.py @@ -98,7 +98,7 @@ def test_perf(self, m, n, dtype, framework, record_property): else: pytest.skip(f"Framework {framework} is not available") with torch.no_grad(): - res = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + res = common.benchmark_framework(framework, framework_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM @@ -213,7 +213,7 @@ def test_perf(self, m, n, dtype, framework, record_property): pytest.skip(f"Framework {framework} is not available") with torch.no_grad(): - res = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + res = common.benchmark_framework(framework, framework_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_linear_gluact_linear.py b/tests/ops/test_linear_gluact_linear.py index f431e006..9cb49653 100644 --- a/tests/ops/test_linear_gluact_linear.py +++ b/tests/ops/test_linear_gluact_linear.py @@ -176,7 +176,7 @@ def test_perf(self, m, n, k, act_type, dtype, backend, arch, record_property): atol=1e-3 if dtype == torch.float16 else 1e-5, ) - res = common.benchmark_framework(backend, backend_fn, use_cudagraph=True) + res = common.benchmark_framework(backend, backend_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_mamba2.py b/tests/ops/test_mamba2.py index 9a9e879c..e820b86d 100644 --- a/tests/ops/test_mamba2.py +++ b/tests/ops/test_mamba2.py @@ -711,7 +711,7 @@ def test_perf(self, T, S, arch, framework, record_property): r["init_state"], chunk_size=chunk_size, ) - result = common.benchmark_framework(framework, framework_fn, use_cupti=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) @@ -785,5 +785,5 @@ def test_perf(self, T, S, arch, framework, record_property): r["dfinal_state"], chunk_size=chunk_size, ) - result = common.benchmark_framework(framework, framework_fn, use_cupti=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) diff --git a/tests/ops/test_matmul.py b/tests/ops/test_matmul.py index 8591f234..4b734b6f 100644 --- a/tests/ops/test_matmul.py +++ b/tests/ops/test_matmul.py @@ -257,7 +257,7 @@ def test_perf( output_processor=output_processor, ) try: - res = common.benchmark_framework(backend, backend_fn, use_cudagraph=False) + res = common.benchmark_framework(backend, backend_fn) except torch.OutOfMemoryError as e: pytest.skip(f"OOM during benchmark: {e}") record_property("benchmark", res) @@ -330,7 +330,7 @@ def test_perf_llm( atol=1e-2, rtol=1e-2, ) - res = common.benchmark_framework(backend, backend_fn, use_cudagraph=False) + res = common.benchmark_framework(backend, backend_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM @@ -567,7 +567,7 @@ def test_perf( pytest.skip(f"Backend {backend} not supported") # Run benchmarks - res = common.benchmark_framework(backend, backend_fn, use_cudagraph=False) + res = common.benchmark_framework(backend, backend_fn) # Record results for reporting record_property("benchmark", res) diff --git a/tests/ops/test_mla.py b/tests/ops/test_mla.py index 42684cc5..448df165 100644 --- a/tests/ops/test_mla.py +++ b/tests/ops/test_mla.py @@ -64,7 +64,7 @@ def reference(q, k, v, qpe, kpe, is_causal, scaling=None): if is_causal: if q.size(-2) > 1: rows, cols = torch.triu_indices(qk.shape[-2], qk.shape[-1], offset=1, device=qk.device) - qk[..., rows, cols] = float("-inf") + qk[..., rows, cols] = qk.new_full((), float("-inf")) # Calculate attention weights m = torch.max(qk, dim=-1)[0] @@ -244,12 +244,7 @@ def test_perf(self, batch, heads, seq_len, d_model, d_pe, dtype, is_causal, back else: pytest.skip(f"Backend {backend} is not available") - # pytorch reference uses dynamic tensor creation (torch.triu_indices) which is - # incompatible with CUDA graph capture — disabling cudagraph for that backend - # to prevent capture_epilogue() being skipped on failure, which would leave the - # default CUDA generator in capturing_=True state and corrupt subsequent tests. - use_cudagraph = backend != "pytorch" - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=use_cudagraph) + result = common.benchmark_framework(backend, backend_fn) record_property("benchmark", result) if dtype == torch.bfloat16: diff --git a/tests/ops/test_mla_decoding.py b/tests/ops/test_mla_decoding.py index c8a4b59f..1b017cc9 100644 --- a/tests/ops/test_mla_decoding.py +++ b/tests/ops/test_mla_decoding.py @@ -228,7 +228,7 @@ def test_perf( pytest.skip(f"Framework {framework} is not available") # Run benchmarks - res = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + res = common.benchmark_framework(framework, framework_fn) # Record results for reporting record_property("benchmark", res) diff --git a/tests/ops/test_mla_decoding_split_kv.py b/tests/ops/test_mla_decoding_split_kv.py index 70f166eb..e03f6cd3 100644 --- a/tests/ops/test_mla_decoding_split_kv.py +++ b/tests/ops/test_mla_decoding_split_kv.py @@ -158,5 +158,5 @@ def test_perf( multiple_outputs=False, ) - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) diff --git a/tests/ops/test_moe.py b/tests/ops/test_moe.py index fa432fdc..f1e72e33 100644 --- a/tests/ops/test_moe.py +++ b/tests/ops/test_moe.py @@ -267,7 +267,7 @@ def backend_fn(): # Run performance benchmark try: - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=True) + result = common.benchmark_framework(backend, backend_fn) record_property("benchmark", result) except Exception as e: pytest.fail(f"Performance benchmark failed for {backend} {call_type}: {e}") diff --git a/tests/ops/test_moe_actgrad_bwd.py b/tests/ops/test_moe_actgrad_bwd.py index 1ebb04ed..ab18fceb 100644 --- a/tests/ops/test_moe_actgrad_bwd.py +++ b/tests/ops/test_moe_actgrad_bwd.py @@ -590,8 +590,6 @@ def test_perf(self, T, H, I, E, K, activation_type, dtype, framework, record_pro framework, framework_fn, mode="forward", - use_cudagraph=True, - use_cupti=False, ) except Exception as e: if _is_cutile_unsupported_error(e): diff --git a/tests/ops/test_recurrent_gated_delta_rule.py b/tests/ops/test_recurrent_gated_delta_rule.py index 39ed88bf..a645ded1 100644 --- a/tests/ops/test_recurrent_gated_delta_rule.py +++ b/tests/ops/test_recurrent_gated_delta_rule.py @@ -201,7 +201,7 @@ def test_perf(self, T, framework, record_property, arch): else: pytest.skip(f"Framework {framework} is not available") - result = common.benchmark_framework(framework, framework_fn, use_cupti=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) del q, k, v, g, beta diff --git a/tests/ops/test_rms_norm.py b/tests/ops/test_rms_norm.py index 4677ab92..7e66ccd3 100644 --- a/tests/ops/test_rms_norm.py +++ b/tests/ops/test_rms_norm.py @@ -149,7 +149,7 @@ def test_perf(self, m, n, mode, dtype, framework, record_property, arch): rtol=0.0, atol=5e-2, ) - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_rope.py b/tests/ops/test_rope.py index f4d5106c..c33b28f1 100644 --- a/tests/ops/test_rope.py +++ b/tests/ops/test_rope.py @@ -309,7 +309,7 @@ def test_perf( torch.testing.assert_close(ref_q, test_q, atol=1e-2, rtol=1e-2) torch.testing.assert_close(ref_k, test_k, atol=1e-2, rtol=1e-2) - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_silu_and_mul.py b/tests/ops/test_silu_and_mul.py index 7a51472f..41298303 100644 --- a/tests/ops/test_silu_and_mul.py +++ b/tests/ops/test_silu_and_mul.py @@ -149,7 +149,7 @@ def test_perf(self, batch_size, seq_len, hidden_size, dtype, framework, record_p else: pytest.skip(f"Framework {framework} is not available") - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=(framework != "pytorch")) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_splitk_reduce.py b/tests/ops/test_splitk_reduce.py index a9980676..6111bb39 100644 --- a/tests/ops/test_splitk_reduce.py +++ b/tests/ops/test_splitk_reduce.py @@ -211,7 +211,7 @@ def test_perf( multiple_outputs=False, ) - result = common.benchmark_framework(framework, framework_fn, use_cudagraph=True) + result = common.benchmark_framework(framework, framework_fn) record_property("benchmark", result) # Explicit cleanup to prevent OOM diff --git a/tests/ops/test_swiglu.py b/tests/ops/test_swiglu.py index 6894d27b..3be831f9 100644 --- a/tests/ops/test_swiglu.py +++ b/tests/ops/test_swiglu.py @@ -233,7 +233,7 @@ def test_perf(self, batch_size, seq_len, hidden_size, intermediate_size, backend pytest.skip(f"Cutile backend not available: {e}") # Run benchmarks - result = common.benchmark_framework(backend, backend_fn, use_cudagraph=False) + result = common.benchmark_framework(backend, backend_fn) # Log results record_property("benchmark", result) diff --git a/tests/ops/test_transpose.py b/tests/ops/test_transpose.py index 78141361..5761cb68 100644 --- a/tests/ops/test_transpose.py +++ b/tests/ops/test_transpose.py @@ -63,7 +63,7 @@ def test_perf(self, m, n, dtype, backend, record_property): else: pytest.skip(f"Backend {backend} not supported") - res = common.benchmark_framework(backend, backend_fn, use_cudagraph=False) + res = common.benchmark_framework(backend, backend_fn) record_property("benchmark", res) # Explicit cleanup to prevent OOM From 1397ef9b118edaf40d00d2552781fbf73aa03654 Mon Sep 17 00:00:00 2001 From: Changhao Wu Date: Thu, 24 Sep 2026 12:02:48 +0800 Subject: [PATCH 2/2] style(tilecpp): reflow the nvcc invocation --- src/tilegym/ops/tilecpp/utils/_cuda_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/tilegym/ops/tilecpp/utils/_cuda_utils.py b/src/tilegym/ops/tilecpp/utils/_cuda_utils.py index d97dfac8..5842c90b 100644 --- a/src/tilegym/ops/tilecpp/utils/_cuda_utils.py +++ b/src/tilegym/ops/tilecpp/utils/_cuda_utils.py @@ -305,7 +305,6 @@ def compile_cuda_to_cubin( for inc_path in all_include_paths: cmd.extend(["-I", str(inc_path)]) - cmd.append(str(cu_file)) logger.debug(f"Compiling CUDA kernel: {' '.join(cmd)}") @@ -315,7 +314,12 @@ def compile_cuda_to_cubin( source_output_path.write_text(wrapper_source) logger.debug(f"Saved source to: {source_output_path}") - result = subprocess.run(cmd, capture_output=True, text=True, check=True) + result = subprocess.run( + cmd, + capture_output=True, + text=True, + check=True, + ) if result.stderr: logger.debug(f"nvcc stderr: {result.stderr}")