Run fn with warmup + repeat, return (last_result, median_ms, std_ms).
(self, fn)
| 185 | return torch.cuda.max_memory_allocated(self.device) / (1024**3) |
| 186 | |
| 187 | def _time_kernel(self, fn) -> tuple[Any, float, float | None]: |
| 188 | """Run fn with warmup + repeat, return (last_result, median_ms, std_ms).""" |
| 189 | result = None |
| 190 | for _ in range(max(0, self.warmup)): |
| 191 | result = fn() |
| 192 | self._sync() |
| 193 | |
| 194 | elapsed: list[float] = [] |
| 195 | self._reset_memory() |
| 196 | for _ in range(max(1, self.repeat)): |
| 197 | if self.device.type == "cuda": |
| 198 | start = torch.cuda.Event(enable_timing=True) |
| 199 | end = torch.cuda.Event(enable_timing=True) |
| 200 | start.record() |
| 201 | result = fn() |
| 202 | end.record() |
| 203 | end.synchronize() |
| 204 | elapsed.append(start.elapsed_time(end)) |
| 205 | else: |
| 206 | t0 = time.perf_counter() |
| 207 | result = fn() |
| 208 | self._sync() |
| 209 | elapsed.append((time.perf_counter() - t0) * 1000.0) |
| 210 | |
| 211 | self._sync() |
| 212 | median_ms = statistics.median(elapsed) |
| 213 | std_ms = statistics.stdev(elapsed) if len(elapsed) > 1 else None |
| 214 | return result, median_ms, std_ms |
| 215 | |
| 216 | def _measure_peak_memory_once(self, fn) -> float: |
| 217 | """Measure peak allocation for one invocation after the configured warmup.""" |