perf: use cudaEvent for precise GPU timing in debug logs

- cudaEvent.elapsed_time gives microsecond precision vs perf_counter
- cudaEvent measures actual GPU execution, not just kernel launch
- falls back to time.perf_counter on CPU-only devices
This commit is contained in:
2026-08-08 13:18:11 +08:00
parent 0b661bae85
commit 9bff61fb91
+13 -2
View File
@@ -26,15 +26,26 @@ logger = logging.getLogger(__name__)
@contextmanager @contextmanager
def timed(label: str, log: Optional[logging.Logger] = None): def timed(label: str, log: Optional[logging.Logger] = None):
"""Wall-clock debug timer, enabled when the logger level is DEBUG or lower.""" """GPU-precise timer via CUDA events; falls back to perf_counter on CPU."""
log = log or logger log = log or logger
if not log.isEnabledFor(logging.DEBUG): if not log.isEnabledFor(logging.DEBUG):
yield yield
return return
use_cuda = torch.cuda.is_available()
if use_cuda:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
else:
tic = time.perf_counter() tic = time.perf_counter()
yield yield
if use_cuda:
end.record()
torch.cuda.synchronize()
elapsed_ms = start.elapsed_time(end)
else:
elapsed_ms = (time.perf_counter() - tic) * 1000 elapsed_ms = (time.perf_counter() - tic) * 1000
log.debug("%s %.1fms", label, elapsed_ms) log.debug("%s %.2fms", label, elapsed_ms)
@dataclass @dataclass