perf: benchmark decode via real inference engine

- route decode benchmark through InferenceEngine generate path
- add enable_cuda_graph toggle to engine, scheduler, and executor
- make benchmark --cuda-graph/--no-cuda-graph control the toggle
- hoist local time imports to module top

astrai/inference/engine.py

@@ -74,6 +74,7 @@ class InferenceEngine:
         max_batch_size: int = 1,
         max_seq_len: Optional[int] = None,
         cache: Optional[PagePool] = None,
+        enable_cuda_graph: bool = True,
     ):
         self.model = model
         self.tokenizer = tokenizer
@@ -83,6 +84,7 @@ class InferenceEngine:
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
             cache=cache,
+            enable_cuda_graph=enable_cuda_graph,
         )
 
         self.scheduler.start()

astrai/inference/runtime/executor.py

@@ -188,6 +188,7 @@ class Executor:
         task_cache: TaskCacheManager,
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
+        enable_cuda_graph: bool = True,
     ):
         self.model = model
         self.kv_cache = kv_cache
@@ -221,7 +222,8 @@ class Executor:
         # Enabled at init-time via _warmup_cuda_graphs for CudaBackend
         # on supported head_dims; left disabled otherwise.
         self._graph_ctx = CudaGraphContext()
-        self._try_enable_cuda_graph()
+        if enable_cuda_graph:
+            self._try_enable_cuda_graph()
 
     def _try_enable_cuda_graph(self):
         if not self._graph_supported:

astrai/inference/scheduler.py

@@ -27,6 +27,7 @@ class InferenceScheduler:
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
         cache: Optional[PagePool] = None,
+        enable_cuda_graph: bool = True,
     ):
         config = model.config
 
@@ -74,6 +75,7 @@ class InferenceScheduler:
             task_cache=self._task_cache,
             device=self.device,
             dtype=self.dtype,
+            enable_cuda_graph=enable_cuda_graph,
         )
 
         self._stop_event = threading.Event()