perf: gather 向量化

- astrai/inference/scheduler.py: add_task 增加 max_seq_len 检查，超限时直接发 STOP 信号终止
- astrai/inference/scheduler.py: _maybe_alloc_page 返回 bool，alloc 失败时标记 ABORTED + 发 STOP
- astrai/inference/scheduler.py: _execute_decode 过滤分配失败任务，避免 page_table 越界
- astrai/inference/scheduler.py: _remove_finished_tasks 清理 ABORTED 任务并释放 pages
- astrai/inference/scheduler.py: _execute_prefill input_mask 改为覆盖全部 prompt_len
- astrai/model/transformer.py: seq_mask is None 分支补全 start_pos + seq_len 列

astrai/inference/cache.py

@@ -22,14 +22,16 @@ def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
 
 
 class PagedCache:
-    """Paged KV cache with page-table-indirected read/write.
+    """Paged KV cache with page-table-indirected read/write and persistent prefix caching.
 
     Combines:
       - Page pool (ref-counted alloc/free via bitmask)
       - KV tensor storage (k_cache, v_cache)
       - Prefix-cache hash lookup (page_content_hash -> physical_page_idx)
+      - LRU eviction for persistent cross-batch prefix caching
 
-    Call :meth:`bind` to obtain a batch view for the attention layers.
+    Pages with recorded hashes persist after refcount reaches 0 (pinned).
+    They are evicted via LRU only when alloc() finds no free pages.
     """
 
     def __init__(
@@ -57,6 +59,24 @@ class PagedCache:
         )
         self._page_to_hash: Dict[int, int] = {}
         self._hash_to_page: Dict[int, int] = {}
+        self._lru: List[int] = []
+        self._pin: List[bool] = [False] * n_pages
+
+    def _touch(self, idx: int) -> None:
+        if self._refs[idx] == 0 and idx in self._lru:
+            self._lru.remove(idx)
+            self._lru.append(idx)
+
+    def _evict_one(self) -> int:
+        while self._lru:
+            idx = self._lru.pop(0)
+            h = self._page_to_hash.pop(idx, None)
+            if h is not None:
+                self._hash_to_page.pop(h, None)
+            self._pin[idx] = False
+            self._refs[idx] = 1
+            return idx
+        return -1
 
     def record_page(
         self, page_idx: int, token_ids: List[int], logical_page_idx: int
@@ -67,6 +87,9 @@ class PagedCache:
             self._hash_to_page.pop(old_h, None)
         self._page_to_hash[page_idx] = h
         self._hash_to_page[h] = page_idx
+        self._pin[page_idx] = True
+        if page_idx in self._lru:
+            self._lru.remove(page_idx)
 
     def lookup_prefix(self, token_ids: List[int]) -> List[int]:
         full_pages = len(token_ids) // self.page_size
@@ -76,20 +99,25 @@ class PagedCache:
             p = self._hash_to_page.get(h)
             if p is None:
                 break
+            self._touch(p)
             hits.append(p)
         return hits
 
     def inc_ref(self, idx: int) -> None:
         self._refs[idx] += 1
+        if self._refs[idx] == 1 and idx in self._lru:
+            self._lru.remove(idx)
 
     def alloc(self) -> int:
-        lsb = self._free_mask & -self._free_mask
-        if lsb == 0:
-            return -1
-        idx = lsb.bit_length() - 1
-        self._free_mask ^= lsb
-        self._refs[idx] = 1
-        return idx
+        if self._free_mask:
+            lsb = self._free_mask & -self._free_mask
+            idx = lsb.bit_length() - 1
+            self._free_mask ^= lsb
+            self._refs[idx] = 1
+            if idx in self._lru:
+                self._lru.remove(idx)
+            return idx
+        return self._evict_one()
 
     def alloc_n(self, n: int) -> List[int]:
         pages = [self.alloc() for _ in range(n)]
@@ -103,10 +131,15 @@ class PagedCache:
     def free(self, idx: int) -> None:
         self._refs[idx] -= 1
         if self._refs[idx] == 0:
-            self._free_mask |= 1 << idx
-            h = self._page_to_hash.pop(idx, None)
-            if h is not None:
-                self._hash_to_page.pop(h, None)
+            h = self._page_to_hash.get(idx)
+            if h is not None and self._pin[idx]:
+                self._lru.append(idx)
+            else:
+                self._free_mask |= 1 << idx
+                h = self._page_to_hash.pop(idx, None)
+                if h is not None:
+                    self._hash_to_page.pop(h, None)
+                self._pin[idx] = False
 
     def bind(self, page_table: Tensor, total_len: int = 0) -> "CacheView":
         return CacheView(self, page_table, total_len)
@@ -137,15 +170,13 @@ class PagedCache:
             written += chunk
 
     def gather(self, layer_id: int, page_table: Tensor) -> Tuple[Tensor, Tensor]:
-        k_parts, v_parts = [], []
-        for pi in range(page_table.size(1)):
-            phys_pages = page_table[:, pi]
-            if not (phys_pages >= 0).any():
-                break
-            k_parts.append(self.k_cache[layer_id, phys_pages])
-            v_parts.append(self.v_cache[layer_id, phys_pages])
-        k = torch.cat(k_parts, dim=1)
-        v = torch.cat(v_parts, dim=1)
+        # page_table: [batch, max_pages] with -1 padding for tasks with fewer pages.
+        # clamp(min=0) maps -1 to page 0 (irrelevant data) — truncated by CacheView total_len.
+        safe = page_table.clamp(min=0)
+        k = self.k_cache[layer_id, safe]
+        v = self.v_cache[layer_id, safe]
+        k = k.flatten(1, 2)
+        v = v.flatten(1, 2)
         return k, v
 
 

astrai/inference/engine.py

@@ -11,7 +11,7 @@ import asyncio
 import gc
 import threading
 from dataclasses import dataclass
-from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -126,7 +126,7 @@ class _Result:
             idx: Index of the generation task this token belongs to.
         """
         with self._cond:
-            self.tokens.append(token)
+            self.tokens.append((idx, token))
             if token is not STOP:
                 self.results[idx] += token
             else:
@@ -136,11 +136,11 @@ class _Result:
                     self._cond.notify_all()
             self._event.set()
 
-    def pop_all(self) -> List[str]:
-        """Returns and clears all accumulated tokens.
+    def pop_all(self) -> List[Tuple[int, str]]:
+        """Returns and clears all accumulated (idx, token) pairs.
 
         Returns:
-            List of token strings since the last call.
+            List of (index, token_string) tuples since the last call.
         """
         with self._cond:
             out = self.tokens.copy()
@@ -238,20 +238,22 @@ class InferenceEngine:
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 50,
-    ) -> Union[Generator[str, None, None], str, List[str]]:
+    ) -> Union[Generator, str, List[str]]:
         """Generates text from a prompt.
 
         Args:
             prompt: Single string or list of strings for batch generation.
-            stream: If True, returns a generator yielding tokens one by one.
+            stream: If True, returns a generator yielding tokens.
             max_tokens: Maximum number of tokens to generate.
             temperature: Sampling temperature.
             top_p: Nucleus sampling probability threshold.
             top_k: Top-k sampling count (0 disables).
 
         Returns:
-            Generator (stream=True), single string (non-stream, single prompt),
-            or list of strings (non-stream, batch prompts).
+            stream=False, single prompt: str
+            stream=False, batch: List[str]
+            stream=True, single prompt: Generator[str, None, None]
+            stream=True, batch: Generator[Tuple[int, str], None, None]
         """
         is_batch = isinstance(prompt, list)
         prompts = prompt if is_batch else [prompt]
@@ -348,49 +350,60 @@ class InferenceEngine:
         temperature: float,
         top_p: float,
         top_k: int,
-    ) -> Generator[str, None, None]:
+    ) -> Generator:
         """Internal streaming generator.
 
         Polls the _Result accumulator in a loop, yielding tokens as they arrive.
-        Cleans up the scheduler task on GeneratorExit.
+        Single prompt yields raw token strings; batch yields (idx, token) tuples.
 
         Args:
-            prompts: List of prompts (only first is used; batch not yet supported).
-            is_batch: If True, raises NotImplementedError.
+            prompts: List of prompts.
+            is_batch: If True, yields (idx, token) tuples; else yields raw tokens.
             max_tokens: Maximum tokens to generate.
             temperature: Sampling temperature.
             top_p: Nucleus sampling threshold.
             top_k: Top-k sampling count.
 
         Yields:
-            Decoded token strings.
+            Single prompt: decoded token strings.
+            Batch: (sequence_index, token_string) tuples.
         """
-        if is_batch:
-            raise NotImplementedError("Batch streaming not yet supported")
+        n = len(prompts)
+        result = _Result(count=n)
+        task_ids = []
 
-        result = _Result()
+        for i, p in enumerate(prompts):
+            task_id = self.scheduler.add_task(
+                prompt=p,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                stream_callback=lambda tok, idx=i: result.append(tok, idx),
+            )
+            task_ids.append(task_id)
 
-        task_id = self.scheduler.add_task(
-            prompt=prompts[0],
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            stream_callback=lambda tok: result.append(tok, 0),
-        )
+        remaining = n
+        finished = [False] * n
 
         def gen():
+            nonlocal remaining
             try:
-                while True:
-                    tokens = result.pop_all()
-                    for token in tokens:
+                while remaining > 0:
+                    items = result.pop_all()
+                    for idx, token in items:
                         if token is STOP:
-                            return
-                        yield token
-                    if not result.wait(timeout=0.05):
-                        pass
+                            if not finished[idx]:
+                                finished[idx] = True
+                                remaining -= 1
+                        else:
+                            yield (idx, token) if is_batch else token
+                    if remaining > 0:
+                        if not result.wait(timeout=0.05):
+                            pass
             finally:
-                self.scheduler.remove_task(task_id)
+                for tid in task_ids:
+                    self.scheduler.remove_task(tid)
 
         return gen()
 

astrai/inference/scheduler.py

@@ -147,6 +147,13 @@ class InferenceScheduler:
         if len(prompt_ids) > self.max_prompt_len:
             prompt_ids = prompt_ids[-self.max_prompt_len :]
 
+        if len(prompt_ids) >= self.max_seq_len:
+            if stream_callback:
+                stream_callback(STOP)
+            return task_id
+
+        max_tokens = min(max_tokens, self.max_seq_len - len(prompt_ids))
+
         task = Task(
             task_id=task_id,
             prompt_ids=prompt_ids,
@@ -189,7 +196,10 @@ class InferenceScheduler:
     def _remove_finished_tasks(self) -> None:
         finished = []
         for task in self.active_tasks:
-            if task.is_finished(self.tokenizer.stop_ids):
+            if task.status == TaskStatus.ABORTED:
+                task.finish_time = time.time()
+                finished.append(task)
+            elif task.is_finished(self.tokenizer.stop_ids):
                 task.status = TaskStatus.FINISHED
                 task.finish_time = time.time()
                 finished.append(task)
@@ -203,7 +213,9 @@ class InferenceScheduler:
                 task._pages_freed = True
 
         self.active_tasks = [
-            t for t in self.active_tasks if t.status != TaskStatus.FINISHED
+            t
+            for t in self.active_tasks
+            if t.status not in (TaskStatus.FINISHED, TaskStatus.ABORTED)
         ]
 
     def _refill_active_batch(self) -> None:
@@ -254,7 +266,9 @@ class InferenceScheduler:
 
         seq_len = prompt_len - start_pos
         input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
-        input_mask = torch.ones(batch_sz, seq_len, dtype=torch.bool, device=self.device)
+        input_mask = torch.ones(
+            batch_sz, prompt_len, dtype=torch.bool, device=self.device
+        )
 
         for i, t in enumerate(tasks):
             input_ids[i] = torch.tensor(
@@ -280,10 +294,21 @@ class InferenceScheduler:
             return
 
         tasks = sorted(tasks, key=lambda t: t.task_id)
-        batch_sz = len(tasks)
 
+        valid: List[Task] = []
         for t in tasks:
-            self._maybe_alloc_page(t, start_pos)
+            if self._maybe_alloc_page(t, start_pos):
+                valid.append(t)
+            else:
+                t.status = TaskStatus.ABORTED
+                if t.stream_callback:
+                    t.stream_callback(STOP)
+
+        if not valid:
+            return
+
+        tasks = valid
+        batch_sz = len(tasks)
 
         input_ids = torch.tensor(
             [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
@@ -334,14 +359,15 @@ class InferenceScheduler:
         rows = [t.page_table + [-1] * (max_pages - t.n_pages) for t in tasks]
         return torch.tensor(rows, dtype=torch.long, device=self.device)
 
-    def _maybe_alloc_page(self, task: Task, pos: int) -> None:
+    def _maybe_alloc_page(self, task: Task, pos: int) -> bool:
         needed = self._n_pages_for(pos + 1)
         while task.n_pages < needed:
             p = self.page_cache.alloc()
             if p < 0:
-                break
+                return False
             task.page_table.append(p)
             task.n_pages += 1
+        return True
 
     def _run_generation_loop(self) -> None:
         try:

astrai/model/transformer.py

@@ -29,7 +29,9 @@ def process_attention_mask(
 
     if seq_mask is None:
         if start_pos != 0:
-            seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
+            seq_mask = torch.ones(
+                (1, start_pos + seq_len), dtype=torch.bool, device=device
+            )
         else:
             return None