SlimInfer

参照 DeepInfer（DeepSeek 开源的生产级推理框架）架构思想实现的迷你 LLM 推理框架。

• 纯 PyTorch 实现，零自定义 CUDA 编译（无 Triton / 无 C++ 扩展），Windows 原生可跑
• 目标硬件：RTX 3060 Laptop 6GB 及同档显卡；适配模型：Qwen2.5-1.5B-Instruct（bf16 权重约 3.1GB，实测可用显存约 5.1GB，放得下）
• 完整实现现代推理引擎的四大件：分页 KV cache + 连续批处理 + chunked prefill + 批量采样，可选 CUDA Graph

DeepInfer 面向 Linux + CUDA 13 + SM100（H100/B200）跑 DeepSeek-V4.1-Flash（MoE/MLA）； SlimInfer 把它的架构思想移植到消费级显卡上的稠密 GQA 模型（Qwen2）。

与 DeepInfer 的架构对应

SlimInfer	DeepInfer 原型	说明
sliminfer/memory/block_manager.py	csrc/engine/block_manager.hpp	分页 KV cache 的 block 分配/释放/抢占
sliminfer/memory/kv_cache.py	tilelang_ops/load_paged_kv_cache.py	KV pool、slot 映射、分页 gather
sliminfer/engine/scheduler.py	csrc/engine/scheduler.hpp	连续批处理 + chunked prefill + 抢占调度
sliminfer/engine/llm_engine.py	engine/llm_engine.py	引擎主循环 schedule→forward→sample
sliminfer/model_executor/model_runner.py	worker/model_runner.py	prefill/decode 执行与元数据构建
sliminfer/model_executor/model_runner.py::CudaGraphRunner	worker/cuda_graph_runner.py	decode CUDA Graph（batch×kv_len 二维 bucket）
sliminfer/model_executor/attention.py	layers/attention_backend.py	prefill: SDPA-flash；decode: paged-gather+SDPA
sliminfer/model_executor/layers.py	layers/{layernorm,rotary_embedding}.py	RMSNorm / RoPE / SiLU，数值对齐 HF
sliminfer/model_executor/sampler.py	layers/sampler.py	temperature/top-k/top-p/重复惩罚/逐序列种子
sliminfer/model_executor/model_loader.py	hugging_face/weights.py	safetensors 分片加载
sliminfer/model_executor/models/qwen2.py	models/deepseek.py	稠密 GQA 对应物（权重名与 HF 完全一致）
sliminfer/entrypoints/llm.py	entrypoints/llm.py	离线推理 API
sliminfer/entrypoints/openai_server.py	（DeepInfer 本版未含 serving）	OpenAI 兼容 API + SSE 流式
sliminfer/sequence.py / sampling_params.py	同名文件	序列状态机 / 采样参数（含 deterministic sampling）

未复刻的部分（单卡消费级场景用不上）：MoE EP/EPLB、MLA/DSpark 投机解码、 PD 分离、SM100 内核（FlashMLA/DeepGEMM/DeepEP）、分布式并行。

安装

# 1) torch 必须装 CUDA 版（PyPI 上 Windows 默认是 CPU 版!）
pip install torch --index-url https://download.pytorch.org/whl/cu126
# 国内网络不通时, 从上海交大镜像手动下载 wheel 后本地安装:
#   https://mirror.sjtu.edu.cn/pytorch-wheels/cu126/
#   选 torch-2.x.x+cu126-cp312-cp312-win_amd64.whl

# 2) 其余依赖
pip install -r requirements.txt

# 3) (可选) 安装为命令行工具
pip install -e .

下载模型（自行操作）

SlimInfer 直接读取本地模型目录（含 config.json 与 *.safetensors）。 以 Qwen2.5-1.5B-Instruct 为例（约 3.1GB）：

# 方式 A: huggingface-cli + hf-mirror 镜像（国内推荐）
set HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download Qwen/Qwen2.5-1.5B-Instruct --local-dir ./models/Qwen2.5-1.5B-Instruct

# 方式 B: modelscope
pip install modelscope
modelscope download --model Qwen/Qwen2.5-1.5B-Instruct --local_dir ./models/Qwen2.5-1.5B-Instruct

快速开始

# 1) 离线批量推理 + 对话
python examples/offline_inference.py --model ./models/Qwen2.5-1.5B-Instruct

# 2) 命令行交互聊天（流式输出）
python -m sliminfer.cli --model ./models/Qwen2.5-1.5B-Instruct

# 3) OpenAI 兼容服务器
python -m sliminfer.openai_server --model ./models/Qwen2.5-1.5B-Instruct --port 8000
curl http://127.0.0.1:8000/v1/chat/completions -H "Content-Type: application/json" -d ^
  "{\"messages\":[{\"role\":\"user\",\"content\":\"你好\"}],\"max_tokens\":128}"

# 4) 基准测试（TTFT / decode 速度 / 吞吐）
python examples/benchmark.py --model ./models/Qwen2.5-1.5B-Instruct --num-prompts 8

Python API：

from sliminfer.entrypoints.llm import LLM
from sliminfer.sampling_params import SamplingParams

llm = LLM("./models/Qwen2.5-1.5B-Instruct")
out = llm.chat([{"role": "user", "content": "介绍一下你自己"}],
               SamplingParams(temperature=0.7, top_p=0.8, top_k=20))
print(out.output_text)

6GB 显存调参

bf16 下 Qwen2.5-1.5B 的显存构成：权重 ≈ 3.1GB + KV cache ≈ 28KB/token + 运行时 ≈ 0.9GB。 默认 kv_budget_mb=None 时自动预算：空闲显存 − 权重占用(按 safetensors 文件大小估算) − 1GB 安全余量， 在 3060 Laptop 6GB 上约得 0.9GB KV ≈ 3.3 万 token；手动调整：

参数	默认	说明
--kv-budget-mb	自动	KV cache 预算。512MB ≈ 1.8 万 token ≈ 8 路 × 2K 上下文
--max-model-len	8192	单序列最大长度（受模型上限 32768 约束）
--max-num-seqs	16	并发序列上限；6GB 建议 8~16
--max-num-batched-tokens	1024	每步 prefill chunk 大小，影响 TTFT 与首字延迟
--enable-cuda-graph	关	decode CUDA Graph；小 batch 有 10%~20% 收益，失败自动回退

显存吃紧时优先调小 --kv-budget-mb；桌面端建议独占模式运行（iGPU 输出时 3060 可用显存约 5.1GB）。

工作原理（两步式分页注意力）

                 ┌──────────── 连续批处理 (每步) ────────────┐
   请求队列 ──►  Scheduler: decode 全体 + 1 个 prefill chunk
                 │  · KV 槽位保障 / 显存不足抢占 (recompute)
                 ▼
   ModelRunner ─► InputMetadata (positions / slot_mapping / gather_slots)
                 │
                 ▼
   每层 Attention:
     1) 写: 本步 K/V --slot_mapping scatter──► KV Pool [blocks,16,heads,dim]
     2) 读: 单chunk prefill ──► 本地 K/V + SDPA(flash, is_causal)
            chunked/decode  ──► gather_slots 聚合分页 KV + bool mask + SDPA
                 ▼
   Sampler (top-k/top-p/penalty, 逐序列种子) ──► 停止检查(EOS/stop/长度) ──► 流式输出

KV pool 尾部预留 1 个 dummy block：所有 padding 的读写都指向它， 配合 attention mask 保证分页聚合的正确性。

测试

pytest tests/ -v          # 全部单测 + 与 HF transformers 的正确性对比
pytest tests/ -v -k gpu   # 仅 GPU 冒烟 (tiny 模型 bf16)

测试使用本地生成的 tiny 随机权重 Qwen2（hidden=64, 2 层, vocab=256）， 不需要下载任何模型；正确性门槛是与 HF transformers eager 实现的 logits 逐元素一致（atol=1e-3, fp32）和贪心生成序列完全一致。

Roadmap

• prefix caching（block 级哈希复用，多轮对话场景收益大）
• INT8 weight-only 量化（权重 3.1GB → 1.6GB，给 KV cache 腾空间）
• INT8 KV cache（长上下文显存减半）
• W4A16（GPTQ/AWQ 格式加载）

License

Apache-2.0（与 DeepInfer 一致）