Mini-vLLM Qwen3 Serving
面向 Qwen3-32B 的轻量级大模型推理引擎优化项目。项目实现了请求级 Continuous Batching、Paged KV Cache、跨请求 Prefix Cache、Triton GQA PagedAttention、Tensor Parallel 和 CUDA Graph,并在双 RTX PRO 6000 Blackwell Server Edition 上完成端到端正确性验证与压力测试。
核心优化
- 重构 Scheduler 与 BlockManager 的 KV Block 分配、共享和回收流程,修复请求准入、引用计数、Prefix 命中及 Decode 跨 Block 扩容边界。
- 实现跨请求连续 Prefix Cache 复用;完整缓存块通过 Prefix Hash 匹配,活跃请求共享物理 KV Block,引用计数保证共享块安全释放。
- 使用 Triton 实现 Split-KV GQA PagedAttention Decode Kernel,通过 Block Table 读取非连续 KV Cache,并以 Online Softmax 两阶段归约提升长上下文并行度。
- 完成 Qwen3-32B Packed Weight 的 TP=2 分片加载,直接在目标 GPU 上构建 BF16 模型,避免 CPU FP32 初始化产生的大额临时内存。
- 修复 BF16 TP Logits Gather、稳态 KV Cache 容量估算和 CUDA Graph Replay 的动态 Batch Padding,补充 Worker 异常回收。
- 建立 Transformers、Mini-vLLM Eager、Mini-vLLM CUDA Graph 三后端输出一致性验收,并提供 Prefix Cache、高并发、长 Decode 和 32K Output 四套压力测试。
实测结果
测试模型为 Qwen3-32B BF16,Tensor Parallel Size 为 2,硬件为 2 × NVIDIA RTX PRO 6000 Blackwell Server Edition(每卡 96 GB)。
| 场景 | 配置 | 结果 |
|---|---|---|
| 输出正确性 | Transformers / Eager / CUDA Graph | Transformers 生成前缀匹配,Eager 与 CUDA Graph 完整 Token 序列一致 |
| Prefix Cache | 8 路并发,约 16K Token 共享前缀,每请求生成 256 Token | 命中率 98.69%;Prefill 31.494 s → 1.070 s,降低 96.60%,加速 29.43× |
| 高并发吞吐 | 16 路并发,每请求约 4K Prompt + 256 Output | Prefill 5,908.11 tokens/s;Decode 聚合 436.42 tokens/s |
| 长 Decode | 2 路并发,每请求约 8K Prompt + 4K Output | 连续生成 8,192 Token;Decode 聚合 69.43 tokens/s |
| 最大输出边界 | 单请求约 8K Prompt + 32K Output | 完成 32,768 Token Decode,总上下文 40,937 Token,无 OOM、地址越界或 Graph Replay 错误 |
Triton Kernel 在 RTX 4070 Laptop GPU、Batch=1、Context=4096、Q/KV Heads=64/8、Head Dim=128、FP16 条件下:
| Kernel | 延迟 |
|---|---|
| 旧 Decode Kernel | 2.431 ms |
| Split-KV PagedAttention | 110.868 μs |
| 加速比 | 21.93× |
服务器原始测试结果保存在 tests/results/qwen3_32b,性能数字均来自对应 JSON 文件。
推理链路
Prompt
→ Tokenizer
→ LLMEngine.add_prompt
→ Sequence
→ Scheduler
→ BlockManager(KV Block / Prefix Cache)
→ ModelRunner(Prefill / Decode)
→ Qwen3Model
→ Triton Attention
→ Sampler
→ Scheduler.postprocess
→ 追加 Token,直到 EOS 或达到长度限制Prefill 会将同一轮多个请求中尚未缓存的 Prompt Token 拼接为变长 Batch,同时准备请求边界、KV 写入 Slot 和 Block Table。Decode 每个请求只输入最后一个 Token,通过 Block Table 访问历史 KV Cache,并把新 K/V 写入当前物理 Block。
目录结构
src/myvllm/
├── engine/ # Engine、Scheduler、Sequence、BlockManager、ModelRunner
├── layers/ # Linear、RoPE、Sampler、Triton Attention 等算子
├── models/ # Qwen3 模型结构
└── utils/ # Attention Context 与权重加载工具
tests/
├── benchmarks/ # Qwen3-32B 正确性验收和四套压力测试
├── results/qwen3_32b/ # 双卡服务器原始 JSON 结果
└── test_*.py # Scheduler、Prefix Cache、Loader、Kernel 和运行时边界测试
main_qwen32.py # Qwen3-32B TP=2 推理入口环境
- Linux / WSL2
- Python 3.11
- PyTorch 2.8+(CUDA 12.x)
- Triton 3.x
- Transformers
- 双 GPU 用于 Qwen3-32B BF16 TP=2 部署
uv python install 3.11
uv sync --locked --no-dev --python 3.11
export PYTHONPATH="$PWD/src"运行测试
本地功能回归:
PYTHONPATH="$PWD/src" python -m pytest -qTriton Decode Kernel 正确性与性能:
PYTHONPATH="$PWD/src" python tests/test_attention_large_scale_decode.py --mode correctness
PYTHONPATH="$PWD/src" python tests/test_attention_large_scale_decode.py --mode benchmarkQwen3-32B 三后端正确性验收:
CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH="$PWD/src" \
python tests/benchmarks/validate_qwen3_32b_correctness.py \
--model /path/to/Qwen3-32B运行四套双卡压力测试:
CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH="$PWD/src" \
python tests/benchmarks/benchmark_qwen3_32b_stress.py \
--model /path/to/Qwen3-32B --suite all也可以通过 --suite prefix-cache、high-concurrency、long-decode 或 max-output 独立运行指定场景。
说明
性能结果与 GPU 型号、软件版本、上下文长度和并发配置相关,仓库中的数字仅对应上述测试环境,不代表所有硬件上的通用性能。
本项目基于 Apache-2.0 许可的社区 Mini-vLLM 实现进行优化与扩展。