From 6e5088cc7da6c4b90f233188ffd7358ec773454b Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 8 Aug 2026 12:48:45 +0800 Subject: [PATCH] refactor: remove prefill from CUDA graph warmup - decode capture works without pre-filled KV values - reduces init time and eliminates unused prefill forward --- astrai/inference/core/executor.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index e0aba1d..ba8a621 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -118,28 +118,13 @@ def _warmup_cuda_graphs( pool.task_free(tid) continue - with ( - torch.inference_mode(), - attn_backend(ATTN_BACKEND.CUDA), - timed(f"warmup prefill b={b}", logger), - ): - kv_cache = pool.bind_tasks(task_ids, ws, start_pos=0) - ids_in = torch.tensor(prompt_tokens, dtype=torch.long, device=dev) - pos_in = torch.arange(prompt_len, device=dev).unsqueeze(0).expand(b, -1) - model( - ids_in, - input_mask=pos_in.unsqueeze(-1) >= torch.arange(prompt_len, device=dev), - kv_cache=kv_cache, - position_ids=pos_in, - ) - with ( torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA), timed(f"warmup decode b={b}", logger), ): for step in range(2): - seq_pos = prompt_len + step + seq_pos = step ws.position_ids[:b] = seq_pos for tid in task_ids: pool.task_extend(tid, seq_pos)