release: v1.3.4

- refactor: 分页 KV cache（PagedCache+CacheView）替换固定 slot，删除 PrefixCache
- refactor: 推理引擎控制逻辑重写，修复连续批处理核心缺陷、线程安全问题
- refactor: KV 缓存槽位下沉到注意力层，移除 _remap_kv / _writeback_kv
- refactor: 统一采样路径为 SamplingPipeline batch tensor，删除 apply_sampling_strategies
- refactor: 设计模式优化 inference 模块导入结构（cache/sampling 独立）
- feat: 推理引擎前缀缓存（KV cache 复用）
- feat: OpenAI 兼容 chat completion API（流式+非流式+usage）
- feat: Anthropic 兼容 /v1/messages API，移除旧版 /generate 端点
- feat: GRPO CLI 接入 + on-policy，OpenAI API top_k 参数化
- feat: Checkpoint 支持 extra 通用扩展数据
- feat: Docker Compose 一键部署（GPU/CPU 双模式）
- feat: GRPO 训练参数补充，批处理训练参数表
- fix: 调度器延迟优化 — 移除 5ms 睡眠，修复 refill 任务丢失
- fix: CLI 参数缺失/重复、device_ids 越界、generate 参数名不一致
- fix: 长对话截断方向错误，保留最新 token 而非最早
- fix: remove_task 未释放 KV cache slot 导致第二轮对话死锁
- fix: KV cache 槽位索引错位、版本校验缺失、注意力掩码
- fix: scheduler 越界 bug，SchedulerCallback 回调阶段修正
- perf: _Result 改用 Condition.wait_for 消除非流式 CPU 空转
- perf: decode 每步张量预分配；input_ids 改用一次构建代替逐元素赋值
- refactor: 移除 device_ids 参数，统一 CUDA_VISIBLE_DEVICES
- docs: 更新文档以匹配分页 KV cache 等代码重构
- docs: 修正多处文档错误、补充训练参数说明

astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.3"
+__version__ = "1.3.4"
 __author__ = "ViperEkura"
 
 from astrai.config import (

astrai/inference/cache.py

@@ -4,7 +4,7 @@ Provides:
   - PagedCache: paged KV cache combining page pool and tensor storage.
 """
 
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 import torch
 from torch import Tensor
@@ -12,12 +12,22 @@ from torch import Tensor
 STOP = object()
 
 
+def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
+    start = page_idx * page_size
+    end = min(start + page_size, len(token_ids))
+    h = 0
+    for i in range(start, end):
+        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
+    return h
+
+
 class PagedCache:
     """Paged KV cache with page-table-indirected read/write.
 
     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)
 
     Call :meth:`bind` to obtain a batch view for the attention layers.
     """
@@ -45,6 +55,32 @@ class PagedCache:
             device=device,
             dtype=dtype,
         )
+        self._page_to_hash: Dict[int, int] = {}
+        self._hash_to_page: Dict[int, int] = {}
+
+    def record_page(
+        self, page_idx: int, token_ids: List[int], logical_page_idx: int
+    ) -> None:
+        h = page_hash(token_ids, logical_page_idx, self.page_size)
+        old_h = self._page_to_hash.pop(page_idx, None)
+        if old_h is not None:
+            self._hash_to_page.pop(old_h, None)
+        self._page_to_hash[page_idx] = h
+        self._hash_to_page[h] = page_idx
+
+    def lookup_prefix(self, token_ids: List[int]) -> List[int]:
+        full_pages = len(token_ids) // self.page_size
+        hits: List[int] = []
+        for i in range(full_pages):
+            h = page_hash(token_ids, i, self.page_size)
+            p = self._hash_to_page.get(h)
+            if p is None:
+                break
+            hits.append(p)
+        return hits
+
+    def inc_ref(self, idx: int) -> None:
+        self._refs[idx] += 1
 
     def alloc(self) -> int:
         lsb = self._free_mask & -self._free_mask
@@ -68,6 +104,9 @@ class PagedCache:
         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)
 
     def bind(self, page_table: Tensor, total_len: int = 0) -> "CacheView":
         return CacheView(self, page_table, total_len)

astrai/inference/engine.py

@@ -97,7 +97,8 @@ class _Result:
     """Thread-safe token accumulator for streaming and non-streaming modes.
 
     Supports multiple concurrent generation tasks with per-index result tracking.
-    Uses a threading.Event for efficient waiting on completion.
+    Uses a threading.Condition for efficient completion notification
+    and a threading.Event for streaming wakeup.
     """
 
     def __init__(self, count: int = 1):
@@ -106,7 +107,7 @@ class _Result:
         Args:
             count: Number of concurrent generation tasks to track.
         """
-        self._lock = threading.Lock()
+        self._cond = threading.Condition()
         self._event = threading.Event()
         self.tokens: List[str] = []
         self.results: List[str] = [""] * count
@@ -124,7 +125,7 @@ class _Result:
             token: The decoded token string, or STOP sentinel.
             idx: Index of the generation task this token belongs to.
         """
-        with self._lock:
+        with self._cond:
             self.tokens.append(token)
             if token is not STOP:
                 self.results[idx] += token
@@ -132,6 +133,7 @@ class _Result:
                 if not self._done[idx]:
                     self._done[idx] = True
                     self._completed += 1
+                    self._cond.notify_all()
             self._event.set()
 
     def pop_all(self) -> List[str]:
@@ -140,7 +142,7 @@ class _Result:
         Returns:
             List of token strings since the last call.
         """
-        with self._lock:
+        with self._cond:
             out = self.tokens.copy()
             self.tokens.clear()
             if not out:
@@ -158,13 +160,22 @@ class _Result:
         """
         return self._event.wait(timeout=timeout)
 
+    def wait_completion(self) -> None:
+        """Blocks until all tasks complete (non-streaming).
+
+        Uses a Condition to sleep efficiently instead of busy-waiting.
+        The calling thread is parked until a STOP signal arrives.
+        """
+        with self._cond:
+            self._cond.wait_for(lambda: self._completed >= self._total)
+
     def get_results(self) -> List[str]:
         """Returns all accumulated results for non-streaming mode.
 
         Returns:
             List of complete generated strings, one per task index.
         """
-        with self._lock:
+        with self._cond:
             return self.results.copy()
 
 
@@ -425,8 +436,7 @@ class InferenceEngine:
             )
             task_ids.append(task_id)
 
-        while result._completed < result._total:
-            result.wait(timeout=1.0)
+        result.wait_completion()
 
         for task_id in task_ids:
             self.scheduler.remove_task(task_id)

astrai/inference/scheduler.py

@@ -5,7 +5,7 @@ import threading
 import time
 import uuid
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -53,6 +53,7 @@ class Task:
         self.output_tokens: int = 0
         self.page_table: List[int] = []
         self.n_pages: int = 0
+        self._prefix_cached_tokens: int = 0
         self.arrival_time = time.time()
         self.finish_time: Optional[float] = None
         self.stream_callback = stream_callback
@@ -88,8 +89,8 @@ class InferenceScheduler:
         max_seq_len: Optional[int] = None,
         max_prompt_len: int = 512,
         page_size: int = 64,
-        device: str = "cuda",
-        dtype: torch.dtype = torch.bfloat16,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
     ):
         config = model.config
 
@@ -180,6 +181,11 @@ class InferenceScheduler:
         for idx in indices:
             self.page_cache.free(idx)
 
+    def _record_page_hashes(self, task: Task, start_logical_page: int = 0) -> None:
+        full_pages = len(task.prompt_ids) // self.page_size
+        for i in range(start_logical_page, full_pages):
+            self.page_cache.record_page(task.page_table[i], task.prompt_ids, i)
+
     def _remove_finished_tasks(self) -> None:
         finished = []
         for task in self.active_tasks:
@@ -214,12 +220,25 @@ class InferenceScheduler:
         failed: List[Task] = []
         for task in to_add:
             prompt_len = len(task.prompt_ids)
-            n_pages = self._n_pages_for(prompt_len)
-            task.page_table = self.page_cache.alloc_n(n_pages)
-            if not task.page_table:
+
+            hit_pages = self.page_cache.lookup_prefix(task.prompt_ids)
+            cached_tokens = len(hit_pages) * self.page_size
+            for p in hit_pages:
+                self.page_cache.inc_ref(p)
+
+            remaining = prompt_len - cached_tokens
+            n_new = self._n_pages_for(remaining) if remaining > 0 else 0
+            new_pages = self.page_cache.alloc_n(n_new) if n_new > 0 else []
+
+            if remaining > 0 and not new_pages:
+                for p in hit_pages:
+                    self.page_cache.free(p)
                 failed.append(task)
                 continue
+
+            task.page_table = hit_pages + new_pages
             task.n_pages = len(task.page_table)
+            task._prefix_cached_tokens = cached_tokens
             task.status = TaskStatus.RUNNING
             self.active_tasks.append(task)
 
@@ -227,42 +246,20 @@ class InferenceScheduler:
             with self._lock:
                 self.waiting_queue[:0] = failed
 
-    def _execute_prefill(self) -> None:
-        to_prefill = [t for t in self.active_tasks if t.output_tokens == 0]
-        if not to_prefill:
-            return
-
-        for t in to_prefill:
-            prompt_len = len(t.prompt_ids)
-            t.input_tokens = prompt_len
-            t.output_tokens = 0
-
-        groups: Dict[int, List[Task]] = {}
-        for t in to_prefill:
-            groups.setdefault(len(t.prompt_ids), []).append(t)
-
-        for prompt_len, group in groups.items():
-            self._execute_prefill_batch(group, prompt_len)
-
-    def _execute_prefill_batch(self, tasks: List[Task], prompt_len: int) -> None:
+    def _execute_prefill(
+        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
+    ) -> None:
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
 
-        input_ids = torch.zeros(
-            batch_sz,
-            prompt_len,
-            dtype=torch.long,
-            device=self.device,
-        )
-        input_mask = torch.ones(
-            batch_sz,
-            prompt_len,
-            dtype=torch.bool,
-            device=self.device,
-        )
+        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)
 
         for i, t in enumerate(tasks):
-            input_ids[i] = torch.tensor(t.prompt_ids, device=self.device)
+            input_ids[i] = torch.tensor(
+                t.prompt_ids[start_pos:prompt_len], device=self.device
+            )
 
         page_tables = self._make_page_table_tensor(tasks)
 
@@ -270,10 +267,14 @@ class InferenceScheduler:
             self.model(
                 input_ids,
                 input_mask=input_mask,
-                start_pos=0,
+                start_pos=start_pos,
                 paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
             )
 
+        start_logical_page = start_pos // self.page_size
+        for t in tasks:
+            self._record_page_hashes(t, start_logical_page=start_logical_page)
+
     def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
         if not tasks:
             return
@@ -284,15 +285,21 @@ class InferenceScheduler:
         for t in tasks:
             self._maybe_alloc_page(t, start_pos)
 
-        input_ids = torch.zeros(batch_sz, dtype=torch.long, device=self.device)
-        for i, t in enumerate(tasks):
-            input_ids[i] = t.output_ids[-1] if t.output_ids else t.prompt_ids[-1]
+        input_ids = torch.tensor(
+            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
+            dtype=torch.long,
+            device=self.device,
+        )
 
         active_mask = torch.ones((batch_sz, 1), dtype=torch.bool, device=self.device)
 
         page_tables = self._make_page_table_tensor(tasks)
         total_len = start_pos + 1
 
+        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
+        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
+        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
+
         with torch.inference_mode():
             outputs = self.model(
                 input_ids.unsqueeze(1),
@@ -304,11 +311,9 @@ class InferenceScheduler:
 
         next_tokens = sample(
             logits,
-            temperature=torch.tensor(
-                [t.temperature for t in tasks], device=logits.device
-            ),
-            top_k=torch.tensor([t.top_k for t in tasks], device=logits.device),
-            top_p=torch.tensor([t.top_p for t in tasks], device=logits.device),
+            temperature=temperatures,
+            top_k=top_ks,
+            top_p=top_ps,
         ).tolist()
 
         for t, ntok in zip(tasks, next_tokens):
@@ -349,7 +354,19 @@ class InferenceScheduler:
                     self._task_event.wait(timeout=1.0)
                     continue
 
-                self._execute_prefill()
+                to_prefill = [t for t in self.active_tasks if t.output_tokens == 0]
+                if to_prefill:
+                    for t in to_prefill:
+                        t.input_tokens = len(t.prompt_ids)
+
+                    groups: Dict[Tuple[int, int], List[Task]] = {}
+                    for t in to_prefill:
+                        key = (len(t.prompt_ids), t._prefix_cached_tokens)
+                        groups.setdefault(key, []).append(t)
+
+                    for (prompt_len, start_pos), group in groups.items():
+                        if start_pos < prompt_len:
+                            self._execute_prefill(group, prompt_len, start_pos)
 
                 pos_groups: Dict[int, List[Task]] = {}
                 for t in self.active_tasks:

tests/inference/test_scheduler_concurrency.py

@@ -5,6 +5,7 @@ import time
 from unittest.mock import MagicMock, patch
 
 import pytest
+import torch
 
 from astrai.inference.scheduler import InferenceScheduler
 
@@ -19,6 +20,9 @@ def mock_model_and_tokenizer():
     mock_model.config.dim = 128
     mock_model.config.n_layers = 2
     mock_model.config.max_len = 100
+    mock_model.parameters.return_value = iter(
+        [MagicMock(dtype=torch.float32, device=torch.device("cpu"))]
+    )
 
     mock_tokenizer = MagicMock()
     mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5]