[cuda] Score autotune picks with one event tick of tolerance, on longer kernels

backends/cuda/tests/autotune_test_utils.py

@@ -14,6 +14,7 @@
 import time
 from typing import Iterator, NamedTuple
 
+import torch
 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
@@ -25,7 +26,28 @@ class Pick(NamedTuple):
     )
     name: str
     candidates: int
-    regret: float  # reference time of the pick / reference time of the best candidate
+    regret: float  # _regret() of the pick's and the best's reference times
+    best_us: float  # the best candidate's reference time
+
+
+# Event timing is quantized: cudaEventElapsedTime steps by 0.512 us on the A10G
+# CI GPU (CUDA documents about 0.5 us), hipEventElapsedTime by about 1 us. For
+# a kernel of a few microseconds one step reads as a large ratio (3.072 us vs
+# 2.560 us is 1.2x) although quantization alone can produce it.
+_EVENT_TICK_MS = 1e-3 if torch.version.hip is not None else 0.512e-3
+_FLOAT32_EPS = 2.0**-23
+
+
+def _regret(pick_ms: float, best_ms: float) -> float:
+    """pick / best, scored with one event tick of tolerance.
+
+    A reading is the difference of two quantized timestamps, so it can fall
+    either side of the true time by about a tick. Crediting the best one extra
+    tick keeps a difference the timer cannot resolve from counting as a worse
+    pick; it is a tolerance, not a bound. The margins absorb the float32
+    rounding of the two event times (at most one float32 epsilon each)."""
+    credited = best_ms * (1.0 + 2 * _FLOAT32_EPS) + 1.001 * _EVENT_TICK_MS
+    return max(pick_ms / credited, 1.0)
 
 
 def _spin(stop_at: float) -> None:
@@ -73,7 +95,7 @@ def record_autotune_picks(
 ) -> 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.
+    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."""
@@ -86,9 +108,9 @@ def record(kind, name, timings, reference):
         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()))
-        )
+        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)
@@ -137,5 +159,6 @@ def summarize(picks: list) -> str:
     }
     return (
         f"{len(picks)} decisions {kinds}: mean regret {statistics.mean(regrets):.3f}x, "
-        f"worst {max(regrets):.3f}x, >1.10x: {sum(r > 1.10 for r in regrets)}"
+        f"worst {max(regrets):.3f}x, >1.10x: {sum(r > 1.10 for r in regrets)}, "
+        f"fastest kernel {min(p.best_us for p in picks):.1f} us"
     )

backends/cuda/tests/test_cuda_graph_timing.py

@@ -10,6 +10,7 @@
 """
 
 import statistics
+import struct
 import unittest
 from unittest import mock
 
@@ -21,6 +22,7 @@
     CUDA_GRAPH_AUTOTUNE_TIMING_COMPILE_SPEC,
     CudaBackend,
 )
+from executorch.backends.cuda.tests import autotune_test_utils
 from executorch.backends.cuda.tests.autotune_test_utils import (
     record_autotune_picks,
     SaturatedCpu,
@@ -156,14 +158,16 @@ def test_compile_spec_turns_it_off(self) -> None:
 # End to end: export while the CPU is saturated, and check every pick.
 # ---------------------------------------------------------------------------
 
-_WIDTHS = (256, 768, 1536)
-_TOKENS = 64
+# Large enough that even the smallest kernels run for many event ticks (one
+# tick is ~0.5 us), so a pick's regret reflects the kernels, not the timer.
+_WIDTHS = (512, 1536, 3072)
+_TOKENS = 4096
 _HEADS = 4
 
 
 class _ManyKernels(nn.Module):
-    """Small, but compiles into many distinct reductions, pointwise kernels,
-    matmul templates and triton::sdpa calls, each of them autotuned."""
+    """Compiles into many distinct reductions, pointwise kernels, matmul
+    templates and triton::sdpa calls, each of them autotuned."""
 
     def __init__(self) -> None:
         super().__init__()
@@ -209,6 +213,48 @@ def _export_and_lower(extra_specs) -> None:
     )
 
 
+def _f32(ms: float) -> float:
+    # Event times reach the autotuner as float32.
+    return struct.unpack("f", struct.pack("f", ms))[0]
+
+
+class RegretTest(unittest.TestCase):
+    """Both backends' ticks are checked here: ROCm CI does not run this file."""
+
+    def _regret_in_ticks(self, tick: float, best: int, pick: int) -> float:
+        with mock.patch.object(autotune_test_utils, "_EVENT_TICK_MS", tick):
+            return autotune_test_utils._regret(_f32(pick * tick), _f32(best * tick))
+
+    def test_one_tick_apart_is_not_a_worse_pick(self) -> None:
+        for tick in (0.512e-3, 1e-3):
+            for best, pick in (
+                (0, 0),
+                (0, 1),
+                (5, 5),
+                (5, 6),
+                (10, 11),
+                (31250, 31251),
+            ):
+                with self.subTest(tick=tick, best=best, pick=pick):
+                    self.assertEqual(self._regret_in_ticks(tick, best, pick), 1.0)
+
+    def test_the_best_is_credited_one_tick(self) -> None:
+        for tick in (0.512e-3, 1e-3):
+            for best, pick in ((0, 2), (5, 7), (5, 8), (10, 12)):
+                with self.subTest(tick=tick, best=best, pick=pick):
+                    self.assertAlmostEqual(
+                        self._regret_in_ticks(tick, best, pick),
+                        pick / (best + 1),
+                        places=2,
+                    )
+
+    def test_two_ticks_apart_count_at_any_scale(self) -> None:
+        for tick in (0.512e-3, 1e-3):
+            for best in (31250, 1_000_000):
+                with self.subTest(tick=tick, best=best):
+                    self.assertGreater(self._regret_in_ticks(tick, best, best + 2), 1.0)
+
+
 class SaturatedCpuExportTest(unittest.TestCase):
     # Picks may differ from the paused-CPU reference within measurement noise.
     MAX_REGRET = 1.20

backends/cuda/tests/test_int4_quantized_gemm.py

@@ -550,7 +550,10 @@ def test_logits_top1_and_cosine(self) -> None:
                 self.assertGreaterEqual(cosine.min().item(), 0.999)
 
 
-_GUARD_SHAPES = ((2048, 2048), (6144, 2048), (2048, 8192), (32000, 2048))
+# Every weight packs at least 6 MiB of INT4 data, so each decode kernel runs
+# for many event ticks (one tick is ~0.5 us) and a pick's regret reflects the
+# kernel rather than the timer.
+_GUARD_SHAPES = ((8192, 2048), (6144, 2048), (2048, 8192), (32000, 2048))
 _GUARD_BUCKETS = (1, 4)