From 9bff61fb9128cdf89a42c57f00df6399d32f2b9c Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 8 Aug 2026 13:18:11 +0800
Subject: [PATCH] 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
---
 astrai/inference/core/executor.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py
index 909e17f..fd8c185 100644
--- a/astrai/inference/core/executor.py
+++ b/astrai/inference/core/executor.py
@@ -26,15 +26,26 @@ logger = logging.getLogger(__name__)
 
 @contextmanager
 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
     if not log.isEnabledFor(logging.DEBUG):
         yield
         return
-    tic = time.perf_counter()
+    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()
     yield
-    elapsed_ms = (time.perf_counter() - tic) * 1000
-    log.debug("%s %.1fms", label, elapsed_ms)
+    if use_cuda:
+        end.record()
+        torch.cuda.synchronize()
+        elapsed_ms = start.elapsed_time(end)
+    else:
+        elapsed_ms = (time.perf_counter() - tic) * 1000
+    log.debug("%s %.2fms", label, elapsed_ms)
 
 
 @dataclass