From ae9fd546ef4bdaf03132bb43c6b4b6144d8ae3a1 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 8 Aug 2026 13:05:46 +0800
Subject: [PATCH] perf: merge prefill warmup into _warmup_cuda_graphs

- 64-token prefill forward triggers cuBLAS auto-tuning at init
- reduces first-chat prefill from ~520ms to ~27ms
- warmup decode also drops from ~215ms to ~71ms
---
 astrai/inference/core/executor.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py
index ba8a621..909e17f 100644
--- a/astrai/inference/core/executor.py
+++ b/astrai/inference/core/executor.py
@@ -90,11 +90,35 @@ def _warmup_cuda_graphs(
     model: AutoModel,
     pool: PagePool,
     ws: InferenceWorkspace,
-    gctx: "CudaGraphContext",
+    gctx: CudaGraphContext,
     max_batch_size: int,
     prompt_len: int = 1,
     device: Optional[str] = None,
 ):
+    dev = device or next(model.parameters()).device
+
+    # Prefill warmup: cuBLAS auto-tunes for the actual prompt-length tensor
+    # shapes on first call (F.linear is the dominant cost).  This also warms
+    # up the CUDA context (driver init) and compiles the graph-capture trace
+    # that follows.  Custom .so kernels do NOT need this — they are pre-built.
+    warmup_len = 64
+    tid = "_warmup_prefill"
+    if pool.task_alloc(tid, list(range(warmup_len))):
+        with (
+            torch.inference_mode(),
+            timed("warmup prefill", logger),
+        ):
+            kv = pool.bind_tasks([tid], ws, start_pos=0)
+            ids_in = torch.arange(warmup_len, device=dev).unsqueeze(0)
+            pos_in = ids_in
+            model(
+                ids_in,
+                input_mask=pos_in.unsqueeze(-1) >= torch.arange(warmup_len, device=dev),
+                kv_cache=kv,
+                position_ids=pos_in,
+            )
+        pool.task_free(tid)
+
     batch_sizes = [1]
     n = 2
     while n <= max_batch_size:
@@ -103,10 +127,8 @@ def _warmup_cuda_graphs(
     if max_batch_size not in batch_sizes:
         batch_sizes.append(max_batch_size)
 
-    dev = device or next(model.parameters()).device
-
     for b in batch_sizes:
-        task_ids = [f"_gr_{b}_{i}" for i in range(b)]
+        task_ids = [f"_warmup_decode_{b}_{i}" for i in range(b)]
         prompt_tokens = [list(range(prompt_len)) for _ in range(b)]
         alloc_ok = True
         for tid, pt in zip(task_ids, prompt_tokens):