diff --git a/backends/cuda/autotune/cuda_graph_timing.py b/backends/cuda/autotune/cuda_graph_timing.py index fef4b7bceb8..66a969ea8f2 100644 --- a/backends/cuda/autotune/cuda_graph_timing.py +++ b/backends/cuda/autotune/cuda_graph_timing.py @@ -77,11 +77,33 @@ def __init__(self) -> None: _STATE = _State() + + +def graph_timing_active() -> bool: + """Whether a ``cuda_graph_timing()`` context is active on this thread.""" + return _STATE.stats is not None + + _INSTALL_LOCK = threading.Lock() _installed_previous: Callable[..., Any] | None = None _install_count = 0 +def time_with_cuda_graph_us_counted(fn: Callable[[], Any]) -> float: + """``time_with_cuda_graph_us`` with the out-of-memory retries and the captured + and fallback counts of the active ``cuda_graph_timing()``.""" + stats = _STATE.stats + if stats is None: + return time_with_cuda_graph_us(fn) + try: + us = _time_capturing_robustly(fn, stats) + except Exception: + stats.fallbacks += 1 + raise + stats.captured += 1 + return us + + def _default_cuda_bench( bench: benchmarking.Benchmarker, fn: Callable[[], Any], diff --git a/backends/cuda/tests/autotune_test_utils.py b/backends/cuda/tests/autotune_test_utils.py index 00263b76ec0..3e49bff3f94 100644 --- a/backends/cuda/tests/autotune_test_utils.py +++ b/backends/cuda/tests/autotune_test_utils.py @@ -15,6 +15,7 @@ from typing import Iterator, NamedTuple import torch +import triton from executorch.backends.cuda.autotune import cuda_graph_timing as cgt from torch._inductor.runtime.triton_heuristics import CachingAutotuner from torch._inductor.select_algorithm import AlgorithmSelectorCache @@ -28,6 +29,7 @@ class Pick(NamedTuple): candidates: int regret: float # _regret() of the pick's and the best's reference times best_us: float # the best candidate's reference time + key: tuple = () # our triton.autotune decisions: the autotuner's cache key # Event timing is quantized: cudaEventElapsedTime steps by 0.512 us on the A10G @@ -89,64 +91,134 @@ def paused(self) -> Iterator[None]: os.kill(w.pid, signal.SIGCONT) +def _append_pick( + picks: list, kind: str, name: str, timings: dict, reference: dict, key: tuple = () +) -> None: + """Record a decision with its regret score.""" + valid = {c: t for c, t in reference.items() if t != float("inf")} + used = {c: t for c, t in timings.items() if c in valid and t != float("inf")} + if len(valid) < 2 or not used: + return + pick = min(used, key=used.get) + best = min(valid.values()) + regret = _regret(valid[pick], best) + picks.append(Pick(kind, str(name), len(valid), regret, best * 1000.0, key)) + + +def _measure_kernel_reference( + cpu: SaturatedCpu, orig_kernel, self, args: tuple, kwargs: dict, timings: dict +) -> dict: + """Measure kernel reference timings with CPU paused.""" + with cpu.paused(), cgt.cuda_graph_timing(): + return { + launcher: min(self.bench(launcher, *args, **kwargs) for _ in range(2)) + for launcher in timings + } + + +def _measure_choice_reference( + cpu: SaturatedCpu, orig_choices, cls, choices, autotune_args, timings: dict +) -> dict: + """Measure choice reference timings with CPU paused.""" + with cpu.paused(), cgt.cuda_graph_timing(): + return { + c: min(cls.benchmark_choice(c, autotune_args) for _ in range(2)) + for c in timings + } + + +def _measure_triton_reference( + cpu: SaturatedCpu, self, args: tuple, kwargs: dict, timings: dict +) -> dict: + """Measure triton autotuner reference timings with CUDA-graph timing.""" + + def graph_timed(kernel_call, quantiles): + ms = cgt.time_with_cuda_graph_us(kernel_call) / 1000.0 + return [ms] * len(quantiles) + + self.nargs = dict(zip(self.arg_names, args)) + self.__dict__["do_bench"] = graph_timed + try: + with cpu.paused(): + return { + c: min(self._bench(*args, config=c, **kwargs)[0] for _ in range(2)) + for c in timings + } + finally: + del self.__dict__["do_bench"] + self.nargs = None + + @contextlib.contextmanager def record_autotune_picks( cpu: SaturatedCpu, picks: list, measure: bool = True ) -> Iterator[None]: - """For each autotune decision taken inside the context, re-time every - candidate with the CPU paused (CUDA-graph timing, best of two) and append - a ``Pick`` scoring the pick against the best candidate (see _regret). + """For each autotune decision taken inside the context (Inductor's, and our + kernels' own triton.autotune), re-time every candidate with the CPU paused + (CUDA-graph timing, best of two) and append a ``Pick`` scoring the pick + against the best candidate (see _regret). With ``measure=False`` nothing is re-timed: each decision is recorded with its candidate count and a regret of 1.0.""" orig_kernel = CachingAutotuner.benchmark_all_configs orig_choices = AlgorithmSelectorCache.benchmark_choices - def record(kind, name, timings, reference): - valid = {c: t for c, t in reference.items() if t != float("inf")} - used = {c: t for c, t in timings.items() if c in valid and t != float("inf")} - if len(valid) < 2 or not used: - return - pick = min(used, key=used.get) - best = min(valid.values()) - regret = _regret(valid[pick], best) - picks.append(Pick(kind, str(name), len(valid), regret, best * 1000.0)) - def kernel_hook(self, *args, **kwargs): timings = orig_kernel(self, *args, **kwargs) reference = timings if measure: - with cpu.paused(), cgt.cuda_graph_timing(): - reference = { - launcher: min( - self.bench(launcher, *args, **kwargs) for _ in range(2) - ) - for launcher in timings - } + reference = _measure_kernel_reference( + cpu, orig_kernel, self, args, kwargs, timings + ) self.reset_to_zero_args(*args, **kwargs) kind = "custom" if self.custom_kernel else "inductor" - record(kind, self.inductor_meta.get("kernel_name"), timings, reference) + _append_pick( + picks, kind, self.inductor_meta.get("kernel_name"), timings, reference + ) return timings def choices_hook(cls, choices, autotune_args, *args, **kwargs): timings = orig_choices.__func__(cls, choices, autotune_args, *args, **kwargs) reference = timings if measure: - with cpu.paused(), cgt.cuda_graph_timing(): - reference = { - c: min(cls.benchmark_choice(c, autotune_args) for _ in range(2)) - for c in timings - } - record("template", type(next(iter(timings))).__name__, timings, reference) + reference = _measure_choice_reference( + cpu, orig_choices, cls, choices, autotune_args, timings + ) + _append_pick( + picks, "template", type(next(iter(timings))).__name__, timings, reference + ) return timings + orig_triton_run = triton.runtime.Autotuner.run + + def triton_hook(self, *args, **kwargs): + tuned = set(self.cache) + timings_before = getattr(self, "configs_timings", None) + result = orig_triton_run(self, *args, **kwargs) + new_keys = set(self.cache) - tuned + if not new_keys or getattr(self, "configs_timings", None) is timings_before: + return result + pick = (id(self), next(iter(new_keys))) + timings = { + c: t[0] if isinstance(t, (list, tuple)) else t + for c, t in self.configs_timings.items() + } + if not measure: + _append_pick(picks, "custom", self.base_fn.__name__, timings, timings, pick) + return result + reference = _measure_triton_reference(cpu, self, args, kwargs, timings) + _append_pick(picks, "custom", self.base_fn.__name__, timings, reference, pick) + return result + CachingAutotuner.benchmark_all_configs = kernel_hook AlgorithmSelectorCache.benchmark_choices = classmethod(choices_hook) + triton.runtime.Autotuner.run = triton_hook try: yield finally: CachingAutotuner.benchmark_all_configs = orig_kernel AlgorithmSelectorCache.benchmark_choices = orig_choices + triton.runtime.Autotuner.run = orig_triton_run def summarize(picks: list) -> str: diff --git a/backends/cuda/tests/test_int4_quantized_gemm.py b/backends/cuda/tests/test_int4_quantized_gemm.py index 03d40cfe345..6ed94749e20 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/backends/cuda/tests/test_int4_quantized_gemm.py @@ -25,6 +25,7 @@ import torch import torch.nn.functional as F import triton +from executorch.backends.cuda.autotune.cuda_graph_timing import cuda_graph_timing from executorch.backends.cuda.autotune.launch_params import ( autotune_launch_params, clear_launch_param_cache, @@ -41,6 +42,7 @@ from executorch.backends.cuda.triton.kernels import ( int4_quantized_gemm as int4_kernel, quantized_gemm_family, + quantized_gemm_utils, ) from executorch.backends.cuda.triton.kernels.int4_quantized_gemm import ( int4_autotune_configs, @@ -48,6 +50,7 @@ SUPPORTED_BUCKETS, ) from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import ( + autotune_do_bench, check_split_k, SPLIT_K_CANDIDATES, ) @@ -667,6 +670,105 @@ def linear(*xs): with self.subTest(bucket=bucket): _check_close(self, actual, expected) + def test_a_failing_cuda_graph_timing_falls_back_and_says_so_once(self) -> None: + from executorch.backends.cuda.triton.kernels import quantized_gemm_utils + + x = torch.randn(1 << 16, device="cuda") + + def kernel_call(): + x.mul_(1.0) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), cuda_graph_timing(), mock.patch.object( + quantized_gemm_utils, + "time_with_cuda_graph_us_counted", + side_effect=RuntimeError("capture failed"), + ), mock.patch.object( + quantized_gemm_utils, "_graph_timing_fallback_warned", False + ), self.assertLogs( + quantized_gemm_utils.logger, "WARNING" + ) as logs: + for _ in range(2): + # The quantiles triton.autotune asks for. + ms = quantized_gemm_utils.autotune_do_bench( + kernel_call, (0.5, 0.2, 0.8) + ) + self.assertGreater(ms[0], 0) + self.assertEqual(len(logs.output), 1) + self.assertIn("capture failed", logs.output[0]) + + def test_a_tuned_shape_traces_only_its_picked_config(self) -> None: + """Once the autotuner has picked a config for a shape, tracing hands + Inductor that config alone, so AOTI does not tune it again.""" + torch._dynamo.reset() + n, k, bucket = 832, 1536, 16 + weights = _packed(n, k, seed=41) + x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda") + autotuner = int4_kernel._W4A16_DIRECT_KERNELS[bucket] + self.assertGreater(len(autotuner.configs), 1) + + def linear(x): + return INT4_QUANTIZED_GEMM.op(bucket)(x, *weights, GROUP_SIZE) + + with torch.no_grad(): + eager = linear(x) # tunes this shape + compiled, code = torch._inductor.utils.run_and_get_code( + torch.compile(linear, fullgraph=True), x + ) + source = "\n".join(code) + name = autotuner.fn.fn.__name__ + self.assertIn(f"def {name}(", source) + block = source[ + source.rindex("user_autotune(", 0, source.index(f"def {name}(")) : + ] + configs = block[: block.index("],")] + self.assertEqual(configs.count("'num_warps'"), 1) + _check_close(self, compiled, eager) + + def test_an_untuned_shape_traces_every_config(self) -> None: + mode = FakeTensorMode() + traced = [] + wrap = quantized_gemm_family.wrap_triton + + def recording(kernel): + traced.append(kernel) + return wrap(kernel) + + with mode: + weights = tuple( + torch.empty_like(t, device="cuda") for t in _packed(704, 1280, seed=43) + ) + x = torch.empty(16, 1280, dtype=torch.bfloat16, device="cuda") + with mock.patch.object(quantized_gemm_family, "wrap_triton", recording): + with mode: + int4_kernel._launch(16, x, *weights, GROUP_SIZE, split_k=1) + self.assertIs(traced[0], int4_kernel._W4A16_DIRECT_KERNELS[16]) + + def test_the_picked_config_kernel_is_one_object_per_config(self) -> None: + n, k, bucket = 832, 1536, 16 + weights = _packed(n, k, seed=41) + x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda") + traced = [] + wrap = quantized_gemm_family.wrap_triton + + def recording(kernel): + traced.append(kernel) + return wrap(kernel) + + # Through the op, whose eager path runs the autotuner itself, so the + # first call tunes and caches its pick (wrap_triton's eager path tunes a + # fresh copy on every call and would never reuse it). + with mock.patch.object(quantized_gemm_family, "wrap_triton", recording): + for _ in range(3): + out = INT4_QUANTIZED_GEMM.op(bucket)(x, *weights, GROUP_SIZE) + autotuner = int4_kernel._W4A16_DIRECT_KERNELS[bucket] + tuned = [kernel for kernel in traced if kernel is not autotuner] + self.assertGreaterEqual(len(tuned), 2) + self.assertTrue(all(kernel is tuned[0] for kernel in tuned)) + self.assertEqual(len(tuned[0].configs), 1) + self.assertIn(tuned[0].configs[0], autotuner.cache.values()) + _check_close(self, out, _dequant_matmul(x, *weights, GROUP_SIZE)) + def test_dynamic_m_compiled_bucket_serves_fewer_rows(self) -> None: """A bucket compiled for a dynamic M <= its size masks the extra rows.""" torch._dynamo.reset() @@ -797,10 +899,79 @@ def forward(self, x1_2048, x4_2048, x1_8192, x4_8192): return tuple(outs) +class AutotuneDoBenchTest(unittest.TestCase): + def _on_stream(self, cuda_stream: int): + return mock.patch.object( + torch.cuda, + "current_stream", + return_value=types.SimpleNamespace(cuda_stream=cuda_stream), + ) + + def test_inside_graph_timing_on_a_side_stream_times_in_a_graph(self) -> None: + kernel_call = mock.Mock() + with self._on_stream(1), mock.patch.object( + quantized_gemm_utils, "graph_timing_active", return_value=True + ), mock.patch.object( + quantized_gemm_utils, "time_with_cuda_graph_us_counted", return_value=1500.0 + ) as graph, mock.patch.object( + triton.testing, "do_bench" + ) as bench: + self.assertEqual(autotune_do_bench(kernel_call, None), 1.5) + graph.assert_called_once_with(kernel_call) + bench.assert_not_called() + + def test_out_of_memory_while_capturing_is_retried_and_counted(self) -> None: + from executorch.backends.cuda.autotune import cuda_graph_timing as cgt_module + + kernel_call = mock.Mock() + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), cuda_graph_timing() as stats, mock.patch.object( + cgt_module, + "time_with_cuda_graph_us", + side_effect=[torch.OutOfMemoryError("capture"), 1500.0], + ): + self.assertEqual(autotune_do_bench(kernel_call, None), 1.5) + self.assertEqual( + (stats.oom_retries, stats.captured, stats.fallbacks), (1, 1, 0) + ) + + def test_outside_graph_timing_uses_triton(self) -> None: + kernel_call = mock.Mock() + with self._on_stream(1), mock.patch.object( + quantized_gemm_utils, "graph_timing_active", return_value=False + ), mock.patch.object( + quantized_gemm_utils, "time_with_cuda_graph_us_counted" + ) as graph, mock.patch.object( + triton.testing, "do_bench", return_value=2.0 + ) as bench: + self.assertEqual(autotune_do_bench(kernel_call, None), 2.0) + graph.assert_not_called() + bench.assert_called_once_with( + kernel_call, + warmup=quantized_gemm_utils.AUTOTUNE_WARMUP_MS, + rep=quantized_gemm_utils.AUTOTUNE_REP_MS, + quantiles=None, + ) + + def test_the_legacy_default_stream_uses_triton(self) -> None: + kernel_call = mock.Mock() + with self._on_stream(0), mock.patch.object( + quantized_gemm_utils, "graph_timing_active", return_value=True + ), mock.patch.object( + quantized_gemm_utils, "time_with_cuda_graph_us_counted" + ) as graph, mock.patch.object( + triton.testing, "do_bench", return_value=2.0 + ) as bench: + autotune_do_bench(kernel_call, None) + graph.assert_not_called() + bench.assert_called_once() + + class Int4AutotunePickTest(unittest.TestCase): - """The config AOTInductor picks for each INT4 kernel while compiling is the + """The config picked for each INT4 kernel while AOTInductor compiles is the fastest one, measured afterwards with the CPU idle, even though the compile - ran with every core busy. + ran with every core busy. Our triton.autotune picks it during the + launch-parameter timing, and AOTInductor traces only that pick. On ROCm the CUDA backend times Triton configs without CUDA graphs (see cuda_graph_autotune_timing in cuda_backend.py), so their ranking under a @@ -816,6 +987,29 @@ def setUpClass(cls) -> None: if not torch.cuda.is_available(): raise unittest.SkipTest("CUDA required") + def setUp(self) -> None: + # Start from untuned kernels, so the compile picks every config itself + # rather than reusing the picks other tests made. + autotuners = [ + kernel + for table in ( + int4_kernel._BUCKET_KERNELS, + int4_kernel._W4A16_DIRECT_KERNELS, + int4_kernel._W4A16_SPLIT_KERNELS, + ) + for kernel in table.values() + ] + saved = [(kernel, dict(kernel.cache)) for kernel in autotuners] + for kernel in autotuners: + kernel.cache.clear() + + def restore() -> None: + for kernel, cache in saved: + kernel.cache.clear() + kernel.cache.update(cache) + + self.addCleanup(restore) + def test_compile_time_picks_are_the_fastest_under_cpu_load(self) -> None: from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner @@ -832,18 +1026,45 @@ def test_compile_time_picks_are_the_fastest_under_cpu_load(self) -> None: program = torch.export.export(module, inputs, strict=True) picks = [] measure_picks = torch.version.hip is None + # The keys whose picked config AOTI traces: the configs it compiles in. + traced = set() + tuned_kernel = quantized_gemm_family._tuned_kernel + + def recording_tuned_kernel(kernel, args, kwargs): + chosen = tuned_kernel(kernel, args, kwargs) + # Traced calls pass fake or functional tensors; the timing's eager + # calls (plain tensors) launch every split's pick and compile nothing in. + tracing = any( + isinstance(a, torch.Tensor) and type(a) is not torch.Tensor + for a in args + ) + if tracing and chosen is not kernel: + traced.add( + ( + id(kernel), + quantized_gemm_family._tuning_key(kernel, args, kwargs), + ) + ) + return chosen + specs = [CudaBackend.generate_method_name_compile_spec("forward")] with torch.compiler.config.patch( force_disable_caches=True - ), SaturatedCpu() as cpu: + ), SaturatedCpu() as cpu, mock.patch.object( + quantized_gemm_family, "_tuned_kernel", recording_tuned_kernel + ): with record_autotune_picks(cpu, picks, measure=measure_picks): to_edge_transform_and_lower( program, partitioner=[CudaPartitioner(specs)], compile_config=EdgeCompileConfig(_check_ir_validity=False), ) - int4 = [p for p in picks if p.kind == "custom" and "int4" in p.name] - print("INT4 picks:", summarize(int4)) + tuned = [p for p in picks if p.kind == "custom" and "int4" in p.name] + print("INT4 picks, every split timed:", summarize(tuned)) + # The decisions whose config was compiled in: one per (shape, bucket), + # for the split the launch-parameter timing chose. + int4 = [p for p in tuned if p.key in traced] + print("INT4 picks compiled in:", summarize(int4)) self.assertEqual(len(int4), len(_GUARD_SHAPES) * len(_GUARD_BUCKETS), picks) # prune_configs_by is applied under AOTI: each kernel was tuned over the # configs whose pipeline stages fit its main loop. diff --git a/backends/cuda/triton/kernels/quantized_gemm_family.py b/backends/cuda/triton/kernels/quantized_gemm_family.py index e647dec808b..ef70ff069e6 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -34,6 +34,7 @@ import sys import threading import typing +import weakref from typing import Callable, Optional, Sequence import torch @@ -265,6 +266,42 @@ def _compile_every_candidate_at_once(): _QUEUING_ONLY.reset(token) +# kernel -> {config key: the kernel restricted to that one config}. One object +# per (kernel, config), so every call site of a tuned shape traces the same +# kernel and PyTorch's per-kernel caches keep working. +_SINGLE_CONFIG_KERNELS: weakref.WeakKeyDictionary[object, dict] = ( + weakref.WeakKeyDictionary() +) + + +def _tuned_kernel(kernel, args: Sequence, kwargs: dict): + """``kernel`` restricted to the config its autotuner already picked for + this call's key, or ``kernel`` itself if it has not tuned the key. + + Traced, an autotuned kernel hands Inductor every config, and AOTI then + compiles and benchmarks all of them again for the same shapes; restricted + to the pick, Inductor compiles that one config and benchmarks nothing. + Called eagerly it launches the same config the autotuner would. + """ + if not isinstance(kernel, triton.runtime.Autotuner) or len(kernel.configs) <= 1: + return kernel + try: + config = kernel.cache.get(_tuning_key(kernel, args, kwargs)) + except TypeError: + # A symbolic size in the key (a dynamic M the kernel tunes on): no + # concrete pick to restrict to. + return kernel + if config is None: + return kernel + config_key = tuple(sorted(config.all_kwargs().items())) + variants = _SINGLE_CONFIG_KERNELS.setdefault(kernel, {}) + tuned = variants.get(config_key) + if tuned is None: + tuned = triton.autotune(configs=[config], key=[])(kernel.fn) + variants[config_key] = tuned + return tuned + + def launch_split_k_gemm( kernel, *, @@ -334,7 +371,7 @@ def grid(meta): # The prepass discards the result, so the output is not written. return out _precompile_untuned_configs(kernel, args, kwargs) - wrap_triton(kernel)[grid](*args, **kwargs) + wrap_triton(_tuned_kernel(kernel, args, kwargs))[grid](*args, **kwargs) if split_k > 1: splitk_reduce(partial, out, block_m, reduce_block_n) return out diff --git a/backends/cuda/triton/kernels/quantized_gemm_utils.py b/backends/cuda/triton/kernels/quantized_gemm_utils.py index 535ce67388c..36dbc6a3645 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_utils.py +++ b/backends/cuda/triton/kernels/quantized_gemm_utils.py @@ -20,12 +20,17 @@ from __future__ import annotations +import logging from typing import Iterable, Optional, Sequence import torch import triton import triton.language as tl import triton.testing +from executorch.backends.cuda.autotune.cuda_graph_timing import ( + graph_timing_active, + time_with_cuda_graph_us_counted, +) from executorch.backends.cuda.autotune.launch_params import InvalidLaunchParam from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.library import wrap_triton @@ -36,6 +41,8 @@ # initialize CUDA. Inductor inlines the value into the kernel source it builds. _IS_HIP = tl.constexpr(torch.version.hip is not None) +logger: logging.Logger = logging.getLogger(__name__) + # Activation quantization granularity: one INT8 scale per K32 block, launched in # K256 tiles. Q8_BLOCK = 32 @@ -353,8 +360,39 @@ def check_split_k(split_k: int, k: int, k_per_split_unit: int = Q8_TILE) -> None AUTOTUNE_REP_MS = 20 +_graph_timing_fallback_warned = False + + +def _warn_graph_timing_fallback(error: RuntimeError) -> None: + global _graph_timing_fallback_warned + if not _graph_timing_fallback_warned: + _graph_timing_fallback_warned = True + logger.warning( + "CUDA-graph timing of a quantized GEMM config failed (%s); timing " + "configs with Triton's event-based do_bench instead, which a busy " + "CPU can skew. Further fallbacks are not reported.", + error, + ) + + def autotune_do_bench(kernel_call, quantiles): - """``do_bench`` for the quantized GEMM kernels' ``triton.autotune``.""" + """``do_bench`` for the quantized GEMM kernels' ``triton.autotune``. + + Inside ``cuda_graph_timing()`` on a capturable stream (the launch-parameter + timing runs on one) each config is timed in a CUDA graph with a cold L2, + like the split candidates and Inductor's own candidates during an export: a + saturated CPU makes event-timed readings of these short kernels bimodal and + their picks up to several times slower (see autotune/cuda_graph_timing.py). + Otherwise, on the legacy default stream, or if the capture fails, it falls + back to Triton's timing. + """ + if graph_timing_active() and torch.cuda.current_stream().cuda_stream != 0: + try: + ms = time_with_cuda_graph_us_counted(kernel_call) / 1000.0 + except RuntimeError as error: + _warn_graph_timing_fallback(error) + else: + return [ms] * len(quantiles) if quantiles else ms return triton.testing.do_bench( kernel_call, warmup=AUTOTUNE_WARMUP_MS,