💡 不依赖 vLLM，不依赖 SGLang。 调度器、分页 KV Cache、FlashAttention-2、CUDA Graph、向量化采样——整条 serving 链路全部自己实现， 每一行代码都能被读懂，每一个优化都能被验证。
| 🧠 设计目标 | 📐 实现规模 | ⚡ 单卡实测 |
|---|---|---|
| 教学级可读性 × 工程级性能 | ~1,500 行 Python（不含第三方 kernel） | Qwen3-0.6B @ 256 并发 1,200+ tok/s |
| 特性 | 说明 |
|---|---|
| 📄 Paged KV Cache | Token 粒度的 KV 池 + 每请求独立页表，PagedAttention 同款设计；分配 O(1)，请求结束的瞬间内存即可复用 |
| 🔄 连续批处理 | Prefill 优先调度 + Token 预算准入控制，运行中的请求绝不饿死，排队请求的 TTFT 最低 |
| ⚡ FlashAttention-2 | Prefill / Decode 双路径 attention，全部走 FlashInfer fa2 kernel |
| 🔥 算子融合 | QKV / gate-up 合并投影、带残差的融合 RMSNorm、融合 SiLU·mul，KV 写入是手写的 Triton kernel |
| 🖼 CUDA Graph | Decode batch 按档位捕获后整图重放，逐 kernel 的 CPU launch 开销直接归零 |
| 🎲 向量化采样 | 每个请求独立的 temperature / top-k / top-p，同一个 batched kernel 里一次算完 |
| 📦 Flat Varlen Prefill | 变长 prompt 打包进单个连续张量（cu_seqlens），一次 kernel 跑完整个 prefill batch |
单卡实测 · Qwen3-0.6B · bf16 ·
--cuda-graph开启
| 场景 | 并发 | 输出长度 | 吞吐 (tok/s) |
|---|---|---|---|
| 连续批处理（CUDA Graph） | 256 | 1024 | 1,200+ |
| 连续批处理（eager） | 256 | 1024 | ~700 |
CUDA Graph 带来的提升在 70% 以上——这正是 decode 阶段 kernel launch 开销的真实体量。 同一套调度代码在 0.6B → 8B 单卡区间内行为一致，直接换路径即可压测更大模型。
python benchmark/bench.py --model Qwen/Qwen3-0.6B --num-seqs 256 --cuda-graph
# [CONTINUOUS_BATCH] 总 token: 262144, 耗时: xxx, 吞吐: 120x.xxtok/s
请求生命周期：
qwen_engine/
├── llm/ # 用户侧 LLM 入口（tokenizer + generate 主循环）
├── engine/ # 引擎：GPU 资源、显存规划、forward 循环
│ ├── engine.py # 模型 + KV 池 + attention 后端 + sampler
│ ├── graph.py # decode 的 CUDA graph 捕获 / 重放
│ └── sample.py # 向量化 batch sampler
├── scheduler/ # 连续批处理
│ ├── scheduler.py # prefill 优先的顶层调度
│ ├── prefill.py # 准入控制、varlen prefill batch
│ ├── decode.py # 运行中请求集合
│ ├── cache.py # KV 页分配 / 空闲链表
│ └── table.py # 请求 slot + token 池
├── attention/ # FlashInfer 后端（prefill / decode / graph 捕获）
├── kvcache/ # 分页 KV 池布局
├── kernel/ # Triton KV scatter 写 kernel
├── layers/ # Linear、RMSNorm、RoPE、embedding、激活
└── models/ # 纯推理 Qwen3 + safetensors 权重加载
benchmark/ # 离线吞吐压测
test/ # 单元测试
- Python 3.10+
- PyTorch 2.0+（CUDA）
flashinfer-python、triton、safetensors、transformers
git clone https://github.com/yourname/qwen-engine.git
cd qwen-engine
pip install -e .python -m qwen_engine --model /path/to/Qwen3-0.6B --prompt "你是谁？"
# 多条 prompt 直接组成 batch
python -m qwen_engine --model /path/to/Qwen3-0.6B \
--prompt "你好" --prompt "写一首诗" --max-tokens 128from qwen_engine import LLM, SamplingParams
llm = LLM("/path/to/Qwen3-0.6B", max_running_reqs=32)
outputs = llm.generate(
["用一句话解释 paged attention。"],
SamplingParams(temperature=0.7, max_tokens=128),
)
print(outputs[0]["text"])模型路径也可以直接传 HuggingFace 仓库名（如 Qwen/Qwen3-0.6B），首次使用自动下载。
python benchmark/bench.py --model Qwen/Qwen3-0.6B --num-seqs 256 --cuda-graphPrefill / Decode 拆分 Prefill 一次处理整个 prompt，是 compute-bound；decode 每步每请求只出一个 token，是 memory-bound。两者用不同的 attention kernel，也绝不混在同一个 batch 里——调度器每轮只跑其中一种，资源利用率最大化。
Paged KV Cache 不给每条序列开一块连续 buffer。所有 KV 住进一个 token 粒度的固定池子，每个请求持有页表把逻辑位置映射到物理 slot。分配是 O(1) 的链表操作，请求结束的瞬间内存就能被复用——这就是 PagedAttention 的精髓，也是吞吐的底座。
Prefill 优先调度 在 token 预算和 table slot 允许的前提下，新请求优先准入，保证排队请求的首 token 延迟；等待队列一清空立刻切回 decode，运行中的请求永远不会被饿死。
CUDA Graph 重放 Decode 的 batch shape 只有少数几个 pad 后的档位（1, 2, 4, 8, 16, ...）。每个档位只捕获一次，之后每步把输入拷进静态 buffer 直接重放整张图——逐 kernel 的 CPU launch 开销直接归零，decode 吞吐提升 70%+。
Qwen3 全系直接可用——架构参数全部从 config.json 自动读取，零配置切换：
| 模型 | 参数量 | 备注 |
|---|---|---|
| Qwen3-0.6B / 1.7B / 4B | 0.6–4B | 单卡，迭代最快 |
| Qwen3-8B | 8B | 单卡压测甜点区 |
| Qwen3-14B / 32B | 14–32B | 需要张量并行（Roadmap 中） |
- Chunked prefill — 长 prompt 切分，消除 prefill 对 decode 的阻塞
- Prefix caching — 公共前缀 KV 复用，多轮对话 / agent 场景直接受益
- 张量并行（NCCL） — 解锁 14B / 32B 单节点部署
- Overlap 调度 — 第二条 stream 预取下一批 metadata，隐藏 CPU 开销
- OpenAI 兼容 HTTP server — 一行命令起服务，无缝替换现有客户端
- Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM), 2023
- FlashInfer — attention 与采样 kernel 库
- Triton — GPU kernel 编程语言
如果这个项目对你有帮助，欢迎 ⭐ Star 支持一下！
