feat: 推理引擎前缀缓存(KV cache 复用)

- cache.py: 新增模块级 page_hash() 多项式滚动哈希函数;PagedCache 新增
  record_page/lookup_prefix/inc_ref,free() 自动清理哈希映射
- scheduler.py: Task 新增 _prefix_cached_tokens;_refill_active_batch 先查
  缓存命中页(inc_ref)再分配剩余页;合并 _execute_prefill 为单一方法,
  按 (prompt_len, start_pos) 分组批量执行全量/部分 prefill;
  _record_page_hashes 注册完整页哈希;修复 device/dtype 默认值从硬编码
  改为 None(自动检测模型设备)
- test: mock model 补充 dtype/device 适配自动检测
This commit is contained in:
2026-05-09 23:53:57 +08:00
parent ca4e6b907c
commit 3583c46b66
3 changed files with 96 additions and 40 deletions
@@ -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]