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:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user