From edea02c611ab12ee6eb5eb812dd0d4d39d3d20e8 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 23:04:47 -0700 Subject: [PATCH 1/5] Update [ghstack-poisoned] --- backends/cuda/tests/autotune_test_utils.py | 49 ++++++- .../cuda/tests/test_int4_quantized_gemm.py | 123 +++++++++++++++++- .../triton/kernels/quantized_gemm_family.py | 38 +++++- .../triton/kernels/quantized_gemm_utils.py | 18 ++- 4 files changed, 216 insertions(+), 12 deletions(-) diff --git a/backends/cuda/tests/autotune_test_utils.py b/backends/cuda/tests/autotune_test_utils.py index a7092cea367..b2de69f15d0 100644 --- a/backends/cuda/tests/autotune_test_utils.py +++ b/backends/cuda/tests/autotune_test_utils.py @@ -14,6 +14,7 @@ import time from typing import Iterator, NamedTuple +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 @@ -26,6 +27,7 @@ class Pick(NamedTuple): name: str candidates: int regret: float # reference time of the pick / reference time of the best candidate + key: tuple = () # our triton.autotune decisions: the autotuner's cache key def _spin(stop_at: float) -> None: @@ -69,20 +71,21 @@ def paused(self) -> Iterator[None]: @contextlib.contextmanager def record_autotune_picks(cpu: SaturatedCpu, picks: list) -> 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`` with what the pick costs against the best candidate.""" + """For each autotune decision taken inside the context (Inductor's, and our + kernels' own triton.autotune), re-time every candidate with the CPU paused + (best of two) and append a ``Pick`` with what the pick costs against the + best candidate.""" orig_kernel = CachingAutotuner.benchmark_all_configs orig_choices = AlgorithmSelectorCache.benchmark_choices - def record(kind, name, timings, reference): + def record(kind, name, timings, reference, key=()): 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) picks.append( - Pick(kind, str(name), len(valid), valid[pick] / min(valid.values())) + Pick(kind, str(name), len(valid), valid[pick] / min(valid.values()), key) ) def kernel_hook(self, *args, **kwargs): @@ -107,13 +110,49 @@ def choices_hook(cls, choices, autotune_args, *args, **kwargs): record("template", type(next(iter(timings))).__name__, timings, reference) return timings + orig_triton_run = triton.runtime.Autotuner.run + + def triton_hook(self, *args, **kwargs): + # Our kernels' configs are picked by their own triton.autotune (during + # the launch-parameter timing), and AOTI traces only the pick. + tuned = set(self.cache) + result = orig_triton_run(self, *args, **kwargs) + new_keys = set(self.cache) - tuned + if not new_keys: + return result + timings = { + c: t[0] if isinstance(t, (list, tuple)) else t + for c, t in self.configs_timings.items() + } + + def graph_timed(kernel_call, quantiles): + # The CUDA-graph timing the Inductor hook above uses for its reference. + 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(): + reference = { + 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 + record("custom", self.base_fn.__name__, timings, reference, next(iter(new_keys))) + 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 86a43e8ee96..e454c653938 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/backends/cuda/tests/test_int4_quantized_gemm.py @@ -586,6 +586,73 @@ def linear(*xs): with self.subTest(bucket=bucket): _check_close(self, actual, expected) + 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) + + with mock.patch.object(quantized_gemm_family, "wrap_triton", recording): + for _ in range(3): + out = int4_kernel._launch(bucket, x, *weights, GROUP_SIZE, split_k=1) + 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() @@ -714,9 +781,10 @@ def forward(self, x1_2048, x4_2048, x1_8192, x4_8192): 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.""" MAX_REGRET = 1.10 MAX_MEAN_REGRET = 1.03 @@ -726,6 +794,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 @@ -741,18 +832,40 @@ def test_compile_time_picks_are_the_fastest_under_cpu_load(self) -> None: with torch.no_grad(): program = torch.export.export(module, inputs, strict=True) picks = [] + # 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(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): 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 db5d0e3a1b0..32193a58cc0 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -257,6 +257,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 or config.pre_hook is not 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, *, @@ -315,7 +351,7 @@ def grid(meta): _queue_untuned_configs(kernel, args, kwargs) 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..2976590ce54 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_utils.py +++ b/backends/cuda/triton/kernels/quantized_gemm_utils.py @@ -26,6 +26,7 @@ import triton import triton.language as tl import triton.testing +from executorch.backends.cuda.autotune.cuda_graph_timing import time_with_cuda_graph_us 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 @@ -354,7 +355,22 @@ def check_split_k(split_k: int, k: int, k_per_split_unit: int = Q8_TILE) -> None 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``. + + 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). On the legacy + default stream, or if the capture fails, it falls back to Triton's timing. + """ + if torch.cuda.current_stream().cuda_stream != 0: + try: + ms = time_with_cuda_graph_us(kernel_call) / 1000.0 + except RuntimeError: + pass + else: + return [ms] * len(quantiles) if quantiles else ms return triton.testing.do_bench( kernel_call, warmup=AUTOTUNE_WARMUP_MS, From 4c451dc75faa5920bb08791a7877db5de8d15a0a Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 23:24:41 -0700 Subject: [PATCH 2/5] Update [ghstack-poisoned] --- backends/cuda/tests/test_int4_quantized_gemm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/backends/cuda/tests/test_int4_quantized_gemm.py b/backends/cuda/tests/test_int4_quantized_gemm.py index e454c653938..1b5ecb1027b 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/backends/cuda/tests/test_int4_quantized_gemm.py @@ -642,9 +642,12 @@ 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_kernel._launch(bucket, x, *weights, GROUP_SIZE, split_k=1) + 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) From c1809a1f6f99899ba9abe36e5c244c4f5f995ee6 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sat, 17 Oct 2026 01:36:15 -0700 Subject: [PATCH 3/5] Update [ghstack-poisoned] --- backends/cuda/tests/autotune_test_utils.py | 111 ++++++++++++--------- 1 file changed, 64 insertions(+), 47 deletions(-) diff --git a/backends/cuda/tests/autotune_test_utils.py b/backends/cuda/tests/autotune_test_utils.py index 4db151639d6..b2979193e84 100644 --- a/backends/cuda/tests/autotune_test_utils.py +++ b/backends/cuda/tests/autotune_test_utils.py @@ -91,6 +91,63 @@ 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 @@ -105,55 +162,31 @@ def record_autotune_picks( orig_kernel = CachingAutotuner.benchmark_all_configs orig_choices = AlgorithmSelectorCache.benchmark_choices - def record(kind, name, timings, reference, key=()): - 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 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): - # Our kernels' configs are picked by their own triton.autotune (during - # the launch-parameter timing), and AOTI traces only the pick. tuned = set(self.cache) timings_before = getattr(self, "configs_timings", None) result = orig_triton_run(self, *args, **kwargs) new_keys = set(self.cache) - tuned - # A single pruned config is cached without a benchmark, which leaves - # configs_timings as it was. if not new_keys or getattr(self, "configs_timings", None) is timings_before: return result pick = (id(self), next(iter(new_keys))) @@ -162,26 +195,10 @@ def triton_hook(self, *args, **kwargs): for c, t in self.configs_timings.items() } if not measure: - record("custom", self.base_fn.__name__, timings, timings, pick) + _append_pick(picks, "custom", self.base_fn.__name__, timings, timings, pick) return result - - def graph_timed(kernel_call, quantiles): - # The CUDA-graph timing the Inductor hook above uses for its reference. - 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(): - reference = { - 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 - record("custom", self.base_fn.__name__, timings, reference, pick) + 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 From 0030b71ededa0f0d796d6dd9d0be58941766bce5 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sat, 17 Oct 2026 02:18:18 -0700 Subject: [PATCH 4/5] Update [ghstack-poisoned] --- backends/cuda/tests/autotune_test_utils.py | 25 +++++++++++++++------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/backends/cuda/tests/autotune_test_utils.py b/backends/cuda/tests/autotune_test_utils.py index b2979193e84..3e49bff3f94 100644 --- a/backends/cuda/tests/autotune_test_utils.py +++ b/backends/cuda/tests/autotune_test_utils.py @@ -91,7 +91,9 @@ 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: +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")} @@ -109,9 +111,7 @@ def _measure_kernel_reference( """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) - ) + launcher: min(self.bench(launcher, *args, **kwargs) for _ in range(2)) for launcher in timings } @@ -131,6 +131,7 @@ 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) @@ -166,18 +167,26 @@ def kernel_hook(self, *args, **kwargs): timings = orig_kernel(self, *args, **kwargs) reference = timings if measure: - reference = _measure_kernel_reference(cpu, orig_kernel, self, args, kwargs, 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" - _append_pick(picks, 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: - reference = _measure_choice_reference(cpu, orig_choices, cls, choices, autotune_args, timings) - _append_pick(picks, "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 From 173c99f8d331ffd18d581f06dcc0dea5f7b5ab44 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sat, 17 Oct 2026 16:56:21 -0700 Subject: [PATCH 5/5] Update [ghstack-poisoned] --- backends/cuda/tests/test_int4_quantized_gemm.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/backends/cuda/tests/test_int4_quantized_gemm.py b/backends/cuda/tests/test_int4_quantized_gemm.py index ff24fb314e4..6ed94749e20 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/backends/cuda/tests/test_int4_quantized_gemm.py @@ -681,7 +681,7 @@ def kernel_call(): stream = torch.cuda.Stream() with torch.cuda.stream(stream), cuda_graph_timing(), mock.patch.object( quantized_gemm_utils, - "time_with_cuda_graph_us", + "time_with_cuda_graph_us_counted", side_effect=RuntimeError("capture failed"), ), mock.patch.object( quantized_gemm_utils, "_graph_timing_fallback_warned", False @@ -912,7 +912,7 @@ def test_inside_graph_timing_on_a_side_stream_times_in_a_graph(self) -> None: 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", return_value=1500.0 + quantized_gemm_utils, "time_with_cuda_graph_us_counted", return_value=1500.0 ) as graph, mock.patch.object( triton.testing, "do_bench" ) as bench: @@ -940,7 +940,7 @@ def test_outside_graph_timing_uses_triton(self) -> None: 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" + quantized_gemm_utils, "time_with_cuda_graph_us_counted" ) as graph, mock.patch.object( triton.testing, "do_bench", return_value=2.0 ) as bench: @@ -958,7 +958,7 @@ def test_the_legacy_default_stream_uses_triton(self) -> None: 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" + quantized_gemm_utils, "time_with_cuda_graph_us_counted" ) as graph, mock.patch.object( triton.testing, "do_bench", return_value=2.0 ) as bench: