Hetero-Paged-Infer

CI License Rust

异构推理系统 — 基于 PagedAttention 和 Continuous Batching 的 CPU-GPU 协同推理引擎。

目录

核心特性

特性 说明
PagedAttention KV Cache 分页式显存管理，按需分配/释放物理块，支持 copy-on-write
Continuous Batching 连续批处理调度器，prefill/decode 分阶段管理，decode 优先调度
内存压力感知 可配置的内存阈值，自动拒绝新请求防止 OOM
CUDA Graph 支持 decode 阶段可捕获 CUDA Graph 加速重复执行
模块化架构 Tokenizer / Scheduler / GPU Executor / KV Cache Manager 均通过 trait 抽象

架构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

┌─────────────────────────────────────────────────────────────────────────┐
│                           InferenceEngine                                │
├─────────────────────────────────────────────────────────────────────────┤
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────────────────┐  │
│  │   Tokenizer  │    │   Scheduler  │    │    KV Cache Manager      │  │
│  │  (编码/解码)  │    │ (prefill/    │    │   (BlockPool/PageTable)  │  │
│  │              │    │   decode)    │    │                          │  │
│  └──────────────┘    └──────┬───────┘    └───────────┬──────────────┘  │
│                             │                        │                  │
│                      ┌──────▼───────┐               │                  │
│                      │ Batch Builder│◄──────────────┘                  │
│                      └──────┬───────┘                                  │
│                             │                                          │
│  ───────────────────────────┼────────────────────────────────────────  │
│                      ┌──────▼───────┐                                  │
│                      │ GPU Executor │                                  │
│                      │  (CUDA/GPU)  │                                  │
│                      └──────┬───────┘                                  │
│                             │                                          │
│                      ┌──────▼───────┐                                  │
│                      │   KV Cache   │                                  │
│                      │ (GPU Memory) │                                  │
│                      └──────────────┘                                  │
└─────────────────────────────────────────────────────────────────────────┘

推理流程

1
2
3

请求输入 → 分词(CPU) → 调度(CPU) → 执行(GPU) → 解码(CPU) → 输出
                              ↓
                        KV Cache 管理器

状态机

1
2

Pending → Prefill → Decode → Completed
                  ↘ Failed

快速开始

环境要求

构建

1
2
3

git clone https://github.com/LessUp/hetero-paged-infer.git
cd hetero-paged-infer
cargo build --release

运行

1
2
3
4
5
6
7
8

# 基本使用
cargo run --release -- --input "你好，世界！" --max-tokens 50

# 使用配置文件
cargo run --release -- --config config.example.json --input "你好"

# 查看帮助
cargo run --release -- --help

测试

1
2
3
4
5
6
7
8

# 运行所有测试
cargo test

# 运行文档测试
cargo test --doc

# 生成文档
cargo doc --open

配置参数

参数 默认值 说明
block_size 16 每个物理块容纳的 token 数
max_num_blocks 1024 最大物理块数量
max_batch_size 32 单次调度最大序列数
max_num_seqs 256 系统最大并发序列数
max_model_len 2048 模型最大上下文长度
max_total_tokens 4096 单批次最大 token 总数
memory_threshold 0.9 内存压力阈值 (0.0-1.0)

API 文档

生成文档

1

cargo doc --open

核心类型

类型 说明
InferenceEngine 推理引擎主编排器
EngineConfig 引擎配置
GenerationParams 生成参数
Request 推理请求
Sequence 活跃序列（含 KV Cache）
Scheduler Continuous Batching 调度器
KVCacheManager KV Cache 管理器
GPUExecutor GPU 执行器接口
Tokenizer 分词器接口

文档导航

工程质量

检查项 状态
选择性导出 ✅ 避免命名空间污染
实例级请求 ID ✅ 避免测试间状态泄漏
分层错误体系 ✅ Memory/Config/Validation/Execution/Scheduler → Engine
属性测试 ✅ proptest 验证不变量
CI ✅ fmt + clippy + test + doc
安全审计 ✅ cargo audit

测试覆盖

类型 数量
单元测试 78
属性测试 15
集成测试 13
文档测试 29
总计 135

当前状态

已实现 ✅

未实现 ❌

链接

许可证

本项目采用 MIT 许可证。

最后更新: 2026-04-16