AI Infra · LLM Serving Systems · Code Walkthrough

剖析 mini-sglang:
从五千行代码到生产级推理引擎

SGLang 团队把近 30 万行的生产代码蒸馏成约 5,000 行的 mini-sglang,性能却与完整版几乎持平。本文以源码为主线,逐站拆解多进程架构、请求状态机、Paged/Radix KV Cache、Chunked Prefill、Overlap Scheduling、CUDA Graph、高性能 Attention/MoE Kernel 与张量并行,并标注每一项设计「是什么、为什么慢、怎么优化、论文出处」。

分享主题 mini-sglang 源码剖析 参考文献 16 篇(含 arXiv / 官方博客) 代码基线 sgl-project/mini-sglang
~5,000 行
Python 代码量,对比 SGLang 近 30 万行(LMSYS 官方博客, 2025)
20,711 tok/s
Qwen3-0.6B 离线吞吐,1×H200,优于 nano-vLLM 的 19,638(LMSYS, 2025)
5,811 tok/s
Qwen3-14B 离线吞吐,1×H200,nano-vLLM 为 5,560(LMSYS, 2025)
≈ SGLang
4×H200、Qwen3-32B、Qwen trace 在线 P90 TTFT/TBT 曲线几乎重合(LMSYS, 2025)
01项目定位:为什么需要 mini-sglang 02进程模型与系统架构 03请求状态机与核心数据结构 04Paged KV Cache 与内存布局 05Radix Cache:前缀复用的核心 06Chunked Prefill:切碎长提示 07Overlap Scheduling:双流水 08CUDA Graph:消除发射开销 09Attention 后端与自研 Kernel 10张量并行 TP 11Benchmark 数据解读 12局限、取舍与学习路径
User API ServerFastAPI / OpenAI Tokenizertext → ids Scheduler ×N每 GPU 一个 EngineModel / Graph KV Poolpaged + radix Detokenizer ids → text stream token 逐块流式返回(SSE)
图0全景地图:一个请求沿「API → Tokenizer → Scheduler → Engine → KV Pool」正向前进,生成的 token 经 Detokenizer 流式返回;本文第 2–10 章依次剖析沿途各站。
01

项目定位:为什么需要 mini-sglang

Motivation · 5k vs 300k LOC

SGLang 是当前最主流的开源 LLM 推理引擎之一,但功能堆叠使其 Python 代码量膨胀到近 30 万行。对学习者来说,通读生产代码几乎不可能;对研究者来说,往复杂框架里注入新逻辑极易破坏隐式不变量(implicit invariants),而从零造引擎又要先花数周补齐前端服务、分词、NCCL 通信等基础设施。mini-sglang 正是为解决这两个痛点而生。

两个目标:教学与研究原型

官方博客明确给出了 mini-sglang 的双重定位:其一,作为教学材料,用约 5,000 行高内聚、全类型标注的代码展示现代 serving 引擎的全部核心部件;其二,作为快速研究原型,让系统研究者「开箱即用地站在 SOTA 基线之上」,只改想改的部分。它最早就是 SGLang 团队内部用来快速验证新系统想法的原型。

KEY TAKEAWAY

mini-sglang 不是玩具:它保留了 Radix Cache、Chunked Prefill、Overlap Scheduling、Tensor Parallelism、JIT CUDA kernel 五大现代优化,在线性能与完整 SGLang 几乎一致,同时代码量压缩约 60 倍。

功能全景:它实现了什么

下表把 mini-sglang 的能力按「推理引擎能力栈」分层。可以看到,除了 Speculative Decoding、Pipeline Parallelism、多模态等少数特性,它与生产引擎的能力边界已非常接近。

层次能力关键实现位置
前端服务OpenAI 兼容 API(/v1/chat/completions)、交互式 Shell、流式 SSEserver/、shell.py
请求治理Continuous Batching、请求状态机、abort、采样参数core.py、scheduler/
KV 内存Paged KV Pool、Radix Cache 前缀复用、Naive 对照实现kvcache/、scheduler/cache.py
调度策略Chunked Prefill、Overlap Scheduling、资源预留防死锁scheduler/prefill.py、scheduler.py
执行优化CUDA Graph、混合 Attention 后端(FA/FI/TRT-LLM)engine/graph.py、attention/
高性能算子JIT store/index kernel、Triton fused MoE、sgl-kernel、PyNCCLkernel/、moe/
分布式单机张量并行 TP(all-reduce / all-gather)distributed/、layers/

同类项目对比:nano-vLLM 与 mini-sglang

社区里与它对标的是 Kedreamix 的 nano-vLLM。两者都主打「可读的迷你引擎」,区别在于:nano-vLLM 更短、更适合入门第一篇;mini-sglang 功能更全,额外实现了 Chunked Prefill、Overlap Scheduling、分阶段异构 Attention 后端、多进程在线服务与 TP,且架构上做了大量面向扩展的抽象(基类接口 + 工厂注册 + 插件化)。官方离线 benchmark 也正是以 nano-vLLM 为基线(见第 11 章)。

学习建议

按「请求的一生」来读码:从 server/launch.py 看进程如何拉起,再到 scheduler/scheduler.py 的主循环,然后顺着 _schedule_next_batch → _prepare_batch → engine.forward_batch 一路跟完,最后再回头精读 kvcache/ 与 kernel/ 这两块最难的部分。

[REF-1]
Mini-SGLang: Efficient Inference Engine in a Nutshell
Ziyi Xu, LMSYS Org · 官方博客 · 2025 · 项目动机、特性与 benchmark 的第一手资料
blog
[REF-2]
SGLang: Efficient Execution of Structured Language Model Programs
Zheng et al. · NeurIPS 2024 · RadixAttention 与结构化生成的原始论文
arXiv:2312.07104
02

进程模型与系统架构

Architecture · ZMQ + NCCL

mini-sglang 沿用了 SGLang 的高层架构:前端 API 进程 + Tokenizer 进程 + Detokenizer 进程 + 每 GPU 一个 Scheduler 进程。进程间控制面消息走 ZeroMQ(ZMQ),GPU 之间的张量数据走 NCCL(经 torch.distributed 或自研 PyNCCL)。

API ServerFastAPI · 入口进程 Tokenizer可多实例 Detokenizer全局 1 个 Scheduler 0 (Rank 0)Engine · Model TP-0 · GPU 0 Scheduler 1 (Rank 1)Engine · Model TP-1 · GPU 1 Scheduler 2 (Rank 2)Engine · Model TP-2 · GPU 2 ZMQ PUSH/PULL ZMQ 结果回流 ZMQ PUB/SUB 广播 Rank 间: NCCL 张量 + gloo 控制
图1进程架构:API 进程无状态转发;Tokenizer/Detokenizer 独立成进程,避免分词拖慢调度;每个 GPU 对应一个 Scheduler(即 TP Rank),Rank 0 负责对外收发并向其他 Rank 广播。

四类进程的职责切分

请求生命周期:八步走完一个请求

对照图 0 与图 1,一个请求的完整生命周期是:

  1. 用户向 API Server 发请求;
  2. API Server 转发给 Tokenizer;
  3. Tokenizer 编码为 id,发给 Scheduler Rank 0;
  4. Rank 0 把请求广播给所有其他 Rank(TP>1 时);
  5. 所有 Rank 各自调度,并触发本地 Engine 计算下一个 token;
  6. Rank 0 收集输出 token,发给 Detokenizer;
  7. Detokenizer 转成文本回给 API Server;
  8. API Server 以 SSE 流式返回用户。
原理补充:为什么分词要独立进程

tokenize/detokenize 是纯 CPU 操作且可能触发慢分词器与字符串拼接;若放在调度线程里,会和 batch 调度争抢 CPU、直接造成 GPU 气泡。独立进程 + ZMQ 队列后,这部分开销天然被隔离,也便于水平扩 Tokenizer 实例。

代码地图:minisgl 包的模块边界

源码位于 python/minisgl,目录即职责:core(Req/Batch/Context)、scheduler(调度循环与各 Manager)、engine(TP worker 运行时)、kvcache(KV 池与前缀缓存)、attention(注意力后端)、kernel(自研 CUDA/Triton 算子与 PyNCCL)、layers(带 TP 的模型积木)、models(模型实现与权重加载分片)、message(ZMQ 消息的自动序列化)、server(CLI 与进程拉起)、tokenizer(分词 worker)、llm(离线 Python 接口)。

注意:全局 Context 是隐式总线

Engine 初始化时会创建唯一的 Context 并注册为全局对象(set_global_ctx)。模型各层在 forward 时通过 get_global_ctx() 取当前 batch、page_table、attention backend、kv pool。读码时要始终意识到这条「隐式总线」,否则很难理解层的输入参数为什么这么少。

03

请求状态机与核心数据结构

Req · Batch · Context

mini-sglang 用三个 dataclass 就描述清了推理运行时的全部状态:Req(单个请求)、Batch(一次前向的请求集合)、Context(进程级全局状态)。理解 Req 上三个长度字段,是理解整个调度器的钥匙。

Req:三个长度刻画请求一生

每个 Req 持有输入 token(CPU tensor)、槽位号 table_idx、采样参数与缓存句柄,其中三个长度字段构成状态机:

由此派生出两个常用量:extend_len = device_len - cached_len(本批次真正要算的 token 数),remain_len = max_device_len - device_len(还剩多少 token 要生成)。每跑完一次前向,complete_one() 把 cached_len 推进到 device_len、device_len 再加 1——decode 阶段每次只前进一个 token。

cached_len device_len max_device_len extend_len remain_len(待生成) 已缓存
图2Req 的长度状态机:cached_len → device_len 之间是本批要计算的 extend_len;device_len → max_device_len 之间是待生成的 remain_len。prefill 后每生成一个 token,整条状态向右推进一步。

Batch:prefill 与 decode 两种相位

Batch 是一次前向的请求集合,用 phase 区分 "prefill" 与 "decode"。调度器负责为它填好 input_ids、positions、out_loc;attention 后端负责挂 attn_metadata。注意 padded_reqs 字段:为了适配 CUDA Graph 的固定批量,decode 批次会用 dummy 请求补齐到捕获过的批量大小,真正参与采样的只有前 size 个。

SamplingParams:贪心与采样

采样参数很克制:temperature、top_k、top_p、ignore_eos、max_tokens。is_greedy 判定条件是 temperature≤0(或 top_k=1)且 top_p=1.0。请求进入时调度器还会用 max_seq_len 钳制 max_tokens,超长请求直接丢弃并告警。

KEY TAKEAWAY

Continuous Batching 的本质就体现在状态机里:每个 Req 独立推进自己的三个长度,调度器每轮把「能凑一批」的请求重新组 Batch——prefill 请求和 decode 请求不再被静态绑定,请求可以在任意轮次加入或离开。

Context:一次只服务一个 Batch

Context 持有 page_size、page_table、attention/moe backend、kv_cache,以及当前正在前向的 batch。它用上下文管理器 forward_batch 保证「同一时刻只有一个活跃 batch」,嵌套使用直接断言失败。模型层所有「当前批次」信息都从这里取。

注意:ChunkedReq 是状态机的特例

被分块的请求使用 ChunkedReq(继承 Req),它重写了 can_decode 使其恒为 False、禁用 append_host,从而保证「所有 chunk 跑完之前绝不进入 decode、也不参与采样」;最后一块能一次算完时才退回普通 Req。详见第 6 章。

04

Paged KV Cache 与内存布局

Paged Allocation · MHA Pool

自回归推理中,每个 token 在每一层都会产生一对 K/V 向量,必须缓存到生成结束。显存占用随并发长度线性增长,且请求的生成长度事先不确定,因此需要一套分页(paged)内存管理。这一思想由 vLLM 的 PagedAttention 推广,mini-sglang 采用了相同的页式抽象,但默认 page_size=1。

KV Pool 的物理布局

MHAKVCache 在初始化时一次性分配一个大 tensor,形状为:

$$(2,\; L,\; P,\; s,\; H_{kv}^{local},\; d)$$

其中 2 对应 K 与 V,L 为层数,P 为页数,s 为 page_size,H_kv 为本地 KV 头数(TP 下切分或复制),d 为 head_dim。每层的 K/V 缓冲可视为 (P·s, H_kv, d) 的 token 行数组,token 的物理位置即行号。

每页 KV 的显存账

Engine 依据加载权重后的剩余显存,反推能分多少页。单 token(全层、K+V)的 KV 占用为:

$$b_{token} = 2 \cdot L \cdot H_{kv}^{local} \cdot d \cdot \text{bytes}$$

单页再乘 page_size。Engine 默认用 90% 的初始空闲显存减去模型权重占用,除以单页大小得到页数;TP 各 Rank 还会 all-reduce 取最小空闲显存,显存不均衡超过 2 GiB 直接报错,防止某个 Rank 先 OOM。

Page Table:逻辑位置到物理位置

每个请求在 page_table 中占一行(由 TableManager 分配 table_idx),行内第 j 个元素记录该请求第 j 个逻辑 token 对应的物理位置。调度器在 allocate_paged 中按页对齐统计每个请求 [cached_len, device_len) 区间需要多少新页,从 free_slots 批量分配后,用 _write_page_table 写回页表。

page_table(每请求一行 · 逻辑→物理) 5 12 7 23 41 … free 8 15 … free req A(table_idx=0) req B(table_idx=1) KV Cache Pool(物理 token 槽) slot 5 slot 7 slot 8 slot 12 slot 23 slot 41 slot 15 … 未分配 / 已缓存页 dummy page越界读的安全垫
图3Paged 布局:page_table 的每行是请求的逻辑→物理映射;物理槽位在 KV Pool 中离散分配,同一请求的 token 不必连续。蓝色为已写入、金色为新分配,dummy page 承接 CUDA Graph padding 的越界访问。

写入路径:store kernel 与 dummy page

注意力后端在每层先调用 store_kv,通过 JIT 编译的 store kernel 把本批新算出的 K/V 按 out_loc(物理位置)散写到池子里——每个 warp 负责一个 token 的整行拷贝,支持按元素宽度特化与 Hopper PDL(Programmatic Dependent Launch)。Engine 还多分配 1 页 dummy page 并让 dummy 请求的页表项全部指向它:当 CUDA Graph 用 dummy 请求补齐批量时,多余的 K/V 写入和注意力读取都落在安全页上,不会污染真实数据或触发越界。

原理补充:页表为什么按 128 字节对齐

page_table 宽度按 32 向上对齐(_align_up_32),且存的是原始物理位置而非页号,保证索引张量的内存对齐,attention kernel 可以直接沿行做向量化 gather。

注意:page_size 并非任意可选

默认 page_size=1;FlashInfer 后端当前也只支持 page_size=1。TRT-LLM 后端仅支持 16/32/64,此时 Engine 会强制把 page_size 改为 64;free_slots 始终按页对齐管理。

05

Radix Cache:前缀复用的核心

Radix Tree · Match / Insert / Evict

在线服务中,大量请求共享前缀:同一套 system prompt、同一篇被反复追问的文档、多轮对话的历史。这些前缀的 KV 如果每个请求都重算一遍,prefill 算力被大量浪费。SGLang 的答案是 RadixAttention:用一棵基数树(radix tree)以 token 序列为键、以 KV 物理位置为值,自动完成最长前缀匹配、插入与驱逐。mini-sglang 完整复刻了这套设计。

root You are a helpfulvalue: [5,12,7] assistant.value: [23,41] Translate tovalue: [8,15] Answer concisely.value: [30,33] French: …value: [44,47] 新请求命中路径(ref_count++)
图4Radix Tree:节点保存一段 token key 与对应 KV 物理位置 value,边由 key 的首个 token(或首个 page)索引。虚线框内是一个新请求的最长匹配路径,沿 root → "You are a helpful" → "assistant." → "Answer concisely." 命中三段,只有未命中部分需要计算。

匹配 match_prefix:沿树走最长前缀

请求入场时,CacheManager 用除最后一个 token 外的输入去匹配(最后一个 token 必然要新算)。_tree_walk 从 root 出发:用 key_fn 取当前位置首个 token(page_size>1 时取前 page_size 个组成 tuple)查孩子;命中孩子后用 fast_compare_key 比较节点 key 与输入,找到第一个不同的位置(这是 C++ 实现,std::mismatch 直接按 int32/int64 字长比较,比 Python 逐元素快得多)。若只匹配了节点的一部分,就调用 split_at 把节点裂成两段,返回可继续挂载的父节点。

引用计数与 lock:防止命中的 KV 被驱逐

匹配返回的物理位置只有在句柄加锁后才安全。lock_handle 沿命中路径把每个节点 ref_count 加 1,并把这些 token 从 evictable_size 划到 protected_size;解锁时逆向归还。PrefillAdder 在估算显存前先 lock 一次、估算后发现空间不足再 unlock——这是一次典型的「先占座再验票」的乐观分配。

插入 insert_prefix:请求完成后回写树

请求在 prefill 完成(非分块)或结束时,由 cache_req 把 [0, cached_len) 的 token 与物理位置回写树。插入长度按 page_size 向下对齐;先走树得到已存在前缀 prefix_len,只把剩余部分作为新节点挂上。源码注释把区间拆得很细:已在树中的旧部分要释放物理页(避免重复占用),新挂部分计入缓存,请求未结束时保留尾部并更新句柄、已结束则连尾部一起释放。

驱逐 evict:LRU 叶子优先、级联向上

物理页不够时,_allocate 触发 evict。规则是:只驱逐 ref_count==0 的叶子节点;所有候选叶子按时间戳入最小堆(时间戳在每次树走访问时刷新,因此近似 LRU,代码 __lt__ 即按 timestamp);弹出一个叶子就回收它的 value、从父节点删除,若父节点因此变成无引用的新叶子,则压入堆中级联回收。root 的 ref_count 恒为 1,永不驱逐。

KEY TAKEAWAY

Radix Cache 把 KV 显存分成三态:Free(未分配)、Evictable(已在树中但无引用,可驱逐)、Protected(有请求正在用)。三态大小由 SizeInfo 实时维护,available_size = free 页 + evictable;驱逐是「按需触发」的 lazy eviction,而不是后台扫描。

Naive Cache:自带的对照组

作为对照,NaivePrefixCache 是空实现:match 永远返回长度 0、insert 不落树、size_info 恒为零。加 --cache naive 即可一键对比「有无前缀复用」的差异,官方在线 benchmark 为了公平对比 SGLang 的 --disable-radix,两边都用了 naive 模式。

注意:mini-sglang 与 SGLang 的一处实现差异

Chunked Prefill 各 chunk 之间的 KV 追踪,mini-sglang 没有走 Radix Tree,而是让所有 chunk 共享同一个 table_idx、直接沿 page_table 同一行累积物理页(见第 6 章)。这是教学版刻意做的简化,少了一层树操作,也更容易读懂。

[REF-3]
SGLang / RadixAttention 原始博客与论文
LMSYS Blog 2024-01-17;论文 Zheng et al., NeurIPS 2024 · 自动维护基数树做 KV 复用,多轮/并行采样场景显著减少重复 prefill
arXiv:2312.07104
06

Chunked Prefill:切碎长提示

Sarathi-Serve · max_extend_tokens

朴素 prefill 一次性处理整条 prompt,带来两个问题:一是显存峰值随序列长度平方膨胀(注意力中间矩阵),长 prompt 直接 OOM;二是调度阻塞,一个长 prefill 跑数秒期间,新来的请求和正在 decode 的请求都得干等,TTFT/TBT 出现毛刺。Sarathi-Serve 提出的 Chunked Prefill 把长 prompt 切成固定预算的小块,分多个批次处理,mini-sglang 默认开启,预算由 max_extend_tokens(默认 8192)控制。

batch t1chunk 1(ChunkedReq) batch t2chunk 2(ChunkedReq) batch t3末块(普通 Req) decode逐 token 读 chunk1 之前的 KV 读 chunk1+2 的 KV 读全部前缀 KV 所有 chunk 共享同一 table_idx,物理页沿 page_table 同一行累积;末块完成前 can_decode 恒为 False
图5Chunked Prefill 时间线:长 prompt 被切成三块,分三个批次执行,后继 chunk 通过 page_table 复用前序 KV;只有末块以普通 Req 跑完,之后才允许进入 decode。

为什么慢/会 OOM:激活显存的平方账

推理显存分三块:模型权重(固定)、KV cache(随总长度线性)、中间激活(随单次计算长度平方)。注意力中 $QK^\top$ 的元素数约为 $n^2$,是激活峰值的主要来源。分块后单次 n 从「整条 prompt」降到「一个 chunk」,峰值随之骤降,且每块算完激活立即释放。

$$\text{activation peak} \propto n_{chunk}^{\,2} \cdot d,\qquad n_{chunk} \le \text{budget}$$

怎么实现:PrefillAdder 与 ChunkedReq

调度时 PrefillAdder 逐个尝试把等待中的请求塞进 token_budget:chunk_size = min(budget, 剩余长度),若 chunk_size 小于剩余长度就构造 ChunkedReq,并把 PendingReq.chunked_req 暂存、排回等待队列头部;下一轮调度检测到 chunked_req 非空,直接复用原 table_idx 与 cache_handle 续跑,cached_len 保持不变、device_len 逐块增长。

KEY TAKEAWAY

ChunkedPrefill 同时解决三件事:把激活峰值钳在预算内(防 OOM)、让长请求不再独占调度轮(decode 请求和其他小请求可以插进 chunk 之间)、为 decode 提供稳定的 TBT。它也是 chunked prefill 与 decode 混批(在完整版 SGLang 中)的基础。

注意:chunk 不是越小越好

官方文档明确提示,把 --max-prefill-length 设得很小(如 128)会显著拖慢性能:prefill 本是 compute-bound、大块才能喂饱 Tensor Core,切得过碎会让 prefill 退化成一连串小 kernel,kernel launch 与 attention 固定开销占比飙升。

[REF-4]
Sarathi-Serve: Taming Throughput-Latency Tradeoff in LLM Inference Serving
Agrawal et al. · OSDI 2024 · 提出 chunked prefill 与 stall-free 调度,平滑 TBT
arXiv:2403.02310
07

Overlap Scheduling:让 CPU 与 GPU 并行

NanoFlow · Dual Stream

一次迭代里 GPU 真正在算的同时,CPU 也没闲着:接收消息、组 batch、匹配 radix、写页表、构造 attention 元数据、处理上一批的 token。若严格串行——CPU 准备完再发 GPU、GPU 跑完 CPU 再处理结果——GPU 会在每批之间出现明显气泡。Overlap Scheduling 的思路来自 NanoFlow:GPU 跑当前批时,CPU 在另一路准备下一批并处理上一批的结果,把 CPU 开销藏进 GPU 计算时间里。

scheduler stream(CPU 侧调度 + 元数据) engine stream(GPU 计算) 准备 batch N 处理 N-1 结果+ 准备 batch N+1 处理 N 结果+ 准备 batch N+2 GPU forward N GPU forward N+1 GPU forward N+2 wait_stream copy_done.synchronize()
图6Overlap 时序:CPU 路与 GPU 路完全错开——上一批 GPU 还在算,CPU 已经在处理它的结果并备好下一批;每批开算前用 wait_stream 等元数据就绪,处理结果前用 event 等 GPU→CPU 拷贝完成。

实现机制:两条 CUDA Stream 与三个同步点

Scheduler 初始化时自带一条调度 stream,Engine 自带一条计算 stream。核心循环 overlap_loop 每轮做三件事:非阻塞地收消息并处理;_schedule_next_batch 在调度 stream 上把下一批的页表、位置、元数据全部备好;然后切到 engine stream,先 wait_stream 等调度 stream 的准备工作完成,再发起异步 forward。forward 返回的 GPU→CPU 拷贝由一个 event 记录,下一轮 _process_last_data 先对 event 调 synchronize,确保 token 已拷回再处理。

同步点 1 · wait_stream

计算流等待调度流:保证本批要读的页表/元数据/H2D 拷贝已经落地,避免读到旧数据(IMA)。

同步点 2 · copy_done event

处理结果前等待 GPU→CPU 的异步拷贝完成,否则 CPU 读到的是未填满的 token buffer。

串行对照:normal_loop 与消融开关

设置环境变量 MINISGL_DISABLE_OVERLAP_SCHEDULING=1 即走 normal_loop:每轮严格「收消息→调度→forward→当场等结果→处理」,Nsight 轨迹上能看到 GPU 被 CPU 开销切成一段段。官方 Nsight 截图显示,开启 overlap 后 GPU kernel 行几乎没有空隙,关闭后则有大量停滞。离线 benchmark 中 mini-sglang 对 nano-vLLM 的优势(第 11 章)主要就来自这一机制。

KEY TAKEAWAY

Overlap 的本质是经典生产者-消费者模型 + 异步多流:让昂贵的瓶颈资源(GPU)永不等待廉价资源(CPU)。同理还有 H2D 传输与计算重叠、PDL 的 kernel 间发射重叠,目标都是掩盖非瓶颈环节的延迟。

工程细节:重叠带来的「双释放」陷阱

重叠后,一个请求可能在「上一批结果处理」与「当前批调度」两个时间面上被同时判定为结束,从而 free 两次。源码用 finished_reqs 集合去重、跳过第二次释放,并把释放操作包进 lazy_free_region 批量合并 free_slots,避免在处理中途修改空闲链表。这类时序 bug 正是 overlap 编程最容易踩的坑。

注意:Overlap 与 TTFT 的权衡

默认调度是「prefill 优先」,且每轮只发一个 forward;高并发下大 batch prefill 会排队,TTFT 可能上升。完整版 SGLang 用更细的预算分配与优先级策略缓解;mini-sglang 保留了主机制但策略从简(代码中留有 TODO: support DECODE first)。

[REF-5]
NanoFlow: Towards Optimal Large Language Model Serving Throughput
Hu et al. · 2024 · 提出 CPU 调度与 GPU 计算重叠的流水线,消除调度气泡
arXiv:2408.12757
[REF-6]
SGLang v0.4: Zero-Overhead Scheduler(官方博客)
LMSYS · 2024-12-04 · overlap scheduler 的工程实现与 Nsight 剖析
blog
08

CUDA Graph:消除 decode 的发射开销

Graph Capture · Replay

Decode 阶段每个请求每步只算一个 token,单步计算量小、kernel 数量却和 prefill 一样多(几十上百层 × 每层多个 kernel)。若每个 kernel 都由 CPU 逐个 launch,kernel launch 开销(每次数微秒)累积起来可能与 GPU 计算本身同量级,小模型、大并发下尤其明显。CUDA Graph 把一整次前向的 kernel 序列录制为一张图,之后一次 replay 即可,CPU 发射开销从「数百次 launch」降到「一次 replay」。

捕获:为一组批量大小预录图

GraphRunner 在启动时为一组批量大小捕获图,默认列表是 [1, 2, 4] + 8,16,…,cuda_graph_max_bs(H200 上 max_bs=256,其他卡 160)。捕获时用 dummy 请求组 decode batch,attention 后端先用 init_capture_graph / prepare_for_capture 准备好「地址固定的 wrapper」,再在 torch.cuda.graph 上下文中跑一次完整前向;各张图共享同一个 graph memory pool 以省显存,捕获顺序从大 bs 到小 bs。

重放:拷贝输入 → replay → 取 logits

decode batch 若 size 不超过最大捕获批量即可用图:pad_batch 先把请求补齐到第一个 ≥ size 的捕获档位;replay 时把真实 input_ids、positions、out_loc 拷进捕获期固定地址的 buffer,attention 后端把 wrapper 换成对应档位的 graph wrapper 并重新 plan,然后一次 g.replay(),最后从固定 logits buffer 切出前 size 行去采样。

pad 到捕获档 拷入固定 buffer graph.replay() 固定 logits 取出 sampler.sample 一次 replay ≈ 数百个 kernel 的一次发射
图7CUDA Graph 重放流程:输入地址在捕获期就已固定,运行时只做少量拷贝与一次 replay;prefill 因形状动态、长度可变不走图,仅 decode 走图。
KEY TAKEAWAY

CUDA Graph 要求「静态」:kernel 序列、输入输出地址、workspace 都必须在捕获期确定。这正是 FlashInfer 提供 CUDAGraphBatchDecodeWrapper(用固定 indptr/indices buffer)的原因——动态规划结果写进固定 buffer,图才能安全重放。

注意:图的显存代价

捕获到 256 档意味着为每档保存一份 kernel 序列与 workspace,启动时还能看到明显的捕获进度条与显存占用。显存紧张时可用 --cuda-graph-max-bs 降低上限;设为 0 则完全关闭,回退到 eager 发射。

09

Attention 后端与高性能 Kernel

FA3 + FlashInfer · JIT · MoE

同样的模型结构,用朴素 PyTorch 算子和用 SOTA kernel,性能可以差出数倍。mini-sglang 在这一层做了两件事:用统一抽象接入三套高性能 attention 后端并支持分阶段混合;用 TVM FFI + JIT 自研了 KV 写入、行索引与 MoE 等关键算子。

混合 Attention:prefill 与 decode 用不同后端

prefill 是 compute-bound,需要大块 GEMM 式 attention 把 Tensor Core 喂满;decode 是 memory-bound,瓶颈在按变长序列 gather KV,需要专门的分页 decode kernel 与跨 KV 并行(FlashDecoding 思路)。一个后端很难两头都最优,因此 mini-sglang 用 HybridBackend 按 batch.phase 转发:Hopper(SM90)上默认 FA3 跑 prefill、FlashInfer 跑 decode;Blackwell(SM100)上默认 TRT-LLM;更老的架构默认 FlashInfer。--attn fa,fi 可手动指定。

AttentionLayer prefill:FA3compute-bound · Tensor Core decode:FlashInfermemory-bound · 分页 gather TRT-LLM fmhaBlackwell 默认 · page 16/32/64
图8混合 Attention:HybridBackend 按相位把请求转发给不同后端;FA3 与 FlashInfer 均为 IO 感知的融合注意力,前者强在 prefill 的 WGMMA 指令利用,后者强在分页 decode 与 CUDA Graph 兼容。

FlashAttention 与 FlashInfer:为什么是它们

朴素 attention 先物化 $n\times n$ 的注意力矩阵再做 softmax,显存读写是 $O(n^2)$;FlashAttention 用分块(tiling)+ online softmax 把 QK、softmax、PV 融合在一块 SRAM 里完成,显存读写降到 $O(n)$,既省显存又更快。FA2 重排了 warp 内并行、减少非 FLOPs;FA3 针对 Hopper 用上 WGMMA/TMA、异步化与 PDL,prefill 再进一步。FlashInfer 则面向 serving:以 block-sparse/分页格式承接异构 KV 存储,提供 load-balanced 的 decode 调度(GQA 比例≥4 时默认启用 Tensor Core),并原生兼容 CUDA Graph。mini-sglang 里 FlashInfer 的 plan 放在 CPU 侧、用 pinned buffer 异步 H2D,也能被 overlap 藏掉。

$$\text{naive: } HBM \sim O(n^2)$$
$$\text{flash: } HBM \sim O(n)$$

TVM FFI:更轻的算子绑定

mini-sglang 的自研算子不走 PyTorch 的 C++ 扩展接口,而采用 TVM FFI 做 Python 绑定与 JIT:FFI 层更轻、调用开销更低,并支持符号化形状/类型匹配(TensorMatcher + SymbolicSize)与运行时编译缓存。算子分两类加载:radix 的 fast_compare_key 是 AOT 编译的 CPU 函数;store/index 是按元素宽度特化、首次调用时 JIT 编译并缓存的 CUDA kernel。

三个自研/集成算子在干什么

实践建议:用 NVTX 看 kernel

每层、每个关键算子都打了 NVTX 标注(@nvtx_annotate,如 "Layer_N")。在 Nsight Systems 里能直接看到层级时间片与 kernel 间隙,判断瓶颈到底在 attention、GEMM 还是调度,是调优第一步。

[REF-7]
FlashAttention-1/2/3
Dao et al. · NeurIPS 2022 / 2023 / 2024 · IO 感知融合注意力,FA3 面向 Hopper 的异步与 WGMMA
arXiv:2407.08608
[REF-8]
FlashInfer: Efficient and Customizable Attention Engine for LLM Serving
Ye et al. · MLSys 2025 · 分页 KV、负载均衡 decode、CUDA Graph 兼容
arXiv:2501.01005
10

张量并行:把一个模型摊到多 GPU

Tensor Parallelism · Megatron 式

单卡放不下大模型、或想要更高吞吐时,需要分布式。mini-sglang 只实现单机张量并行(TP):--tp n 启动 n 个 Scheduler 进程、每进程绑一张 GPU,按 Megatron-LM 的方式把权重沿头/隐藏维切开,每层用 all-reduce/all-gather 拼回结果。

QKV / Attention 按 head 切分KV 头均分/复制 MLP 按 hidden 切各算各的部分和 all-reduce(sum)NVLink Embedding / LM Head 按 vocab 切各持一段词表 all-gather 拼接或各算 logits 分片
图9Megatron 式 TP:注意力按头切(GQA 下 KV 头允许复制)、MLP 按隐藏维切,列/行分支后用 all-reduce 归并;Embedding/LM Head 按词表切并 all-gather。前提是 GPU 间有高速 NVLink。

层的 TP:切在哪里、通信在哪

VocabParallelEmbedding 与 ParallelLMHead 按词表维切分;Attention 的 Q/O 头按 query 头均分,KV 头用 div_even(..., allow_replicate=True)——KV 头数不能被 TP 整除时直接复制(GQA 的常见处理);MLP 的 gate/up/down 按隐藏维切分。按 Megatron 的行列并行结论:行分支(如 down/o_proj)之后各卡持有部分和,一次 all-reduce SUM 即得完整结果。

数据面:NCCL 与自研 PyNCCL

默认 use_pynccl=True:CPU 控制面用 gloo 进程组,GPU 张量通信走 mini-sglang 自研的 PyNCCL communicator(CUDA 实现,初始化时按最大 forward 长度×hidden 预算通信 buffer),all_reduce/all_gather 直接在自定义 NCCL kernel 里完成;关闭 PyNCCL 则回退 torch.distributed 的 NCCL 后端。TP 各 Rank 每轮调度必须看到完全相同的请求序列,否则会因 collective 顺序不一致而 hang。

控制面:gloo 报数 + ZMQ 广播

Rank 0 收到请求后,先用 gloo broadcast 一个「本批消息条数」张量,其他 Rank 据此知道要收几条;再用 ZMQ PUB/SUB 把序列化后的消息原文广播出去,各 Rank 解码后得到一致输入。结果回传只由 Rank 0 做,其他 Rank 静默。把「条数」与「内容」分两路,是为了让非 0 Rank 能精确对齐 collective 次数。

注意:TP 不是线性加速

每层至少一次 all-reduce,通信量随 hidden 规模增长;没有 NVLink 时通信可能吃掉并行收益,这也是 mini-sglang 只支持单机 TP、不做跨机 TP/PP 的原因。启动时各 Rank 还会校验显存差不超过 2 GiB。

[REF-9]
Megatron-LM: Training Multi-Billion Parameter Language Models Using Tensor Parallelism
Shoeybi et al. · 2019 · 列/行切分 + 前向 F、反向 G 的 TP 通信范式
arXiv:1909.08053
[REF-10]
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
Ainslie et al. · EMNLP 2023 · KV 头分组,TP 下 KV 头复制/切分的依据
arXiv:2305.13245
11

Benchmark:五千行为什么能打

Offline + Online · H200

官方给出了离线吞吐与在线延迟两组实验,结论很有说服力:离线小幅领先 nano-vLLM,在线与完整 SGLang 几乎重合。下面逐组拆解数据与口径。

离线吞吐:稳定领先 nano-vLLM

配置:1×H200;Qwen3-0.6B 与 Qwen3-14B;256 条序列;输入/输出长度均在 100–1024 间随机采样;指标为总吞吐(tokens/s)。

05k10k15k20k 19,638 20,711 5,560 5,811 Qwen3-0.6BQwen3-14B nano-vLLM mini-sglang throughput (tokens/s)
图10离线吞吐对比(数据:LMSYS 官方博客,1×H200,256 条长度随机序列):mini-sglang 在两个模型上分别为 20,711 / 5,811 tokens/s,领先 nano-vLLM 约 5.5% / 4.5%,差距主要来自 overlap scheduling。
模型nano-vLLMmini-sglang相对提升
Qwen3-0.6B19,638 tok/s20,711 tok/s+5.5%
Qwen3-14B5,560 tok/s5,811 tok/s+4.5%

小模型(0.6B)下单步 GPU 计算极短、CPU 开销占比更高,因此 overlap 的收益更明显;这也解释了为什么提升幅度在小模型上略大。可用 MINISGL_DISABLE_OVERLAP_SCHEDULING=1 自行做消融。

在线延迟:与完整 SGLang 几乎重合

配置:4×H200(NVLink 互联);Qwen3-32B,TP=4;回放 Qwen trace 的前 1,000 条真实请求;画吞吐–P90 TTFT 与吞吐–P90 TBT 曲线。为公平起见,两边都关闭 radix(mini-sglang 用 --cache naive,SGLang 用 --disable-radix,decode attention 均为 FlashInfer)。

KEY TAKEAWAY

两条曲线在整个吞吐区间几乎重合:P90 TTFT 约 230–470 ms、P90 TBT 约 10–40 ms 的区间内,mini-sglang 与 SGLang 没有系统性差距。这证明在线服务的性能主体由「调度机制 + 底层 kernel」决定,而这两块 mini-sglang 都已完整保留;被砍掉的是功能广度,不是性能。

复现实验入口

注意口径

上述数字均来自官方在 H200/Hopper 环境的实测,换硬件(尤其消费级卡、无 NVLink)后绝对值与最优后端组合都会变化;引用这些数字时应同时带上硬件、模型与数据集口径。

12

局限、取舍与学习路径

Tradeoffs · Reading Guide

mini-sglang 的「小」是主动取舍的结果。读懂它放弃了什么,和读懂它实现了什么同样重要。

功能边界:这些它没有做

维度mini-sglang完整 SGLang
并行策略仅单机 TPTP / PP / DP-Serve / EP / 跨机
多级 KV仅 GPU 显存HiCache(CPU/SSD 分层)、PD 分离
解码加速无Speculative Decoding、EAGLE、Medusa 等
模型/模态Llama / Qwen2/3(含 MoE),纯文本数十种架构、多模态、Embedding/重排序
调度策略prefill 优先、策略单一多种队列/预算/优先级与混批策略
平台仅 Linux(CUDA 依赖)Linux,ROCm 等更多后端

三处关键取舍

推荐阅读顺序(约四遍)

第一遍跑通:按 README 起服务、玩 Shell、跑 benchmark;第二遍跟主循环:launch.py → scheduler.overlap_loop → _prepare_batch → engine.forward_batch;第三遍啃内存:kvcache(pool/naive/radix)→ scheduler/cache.py → kernel/radix、store、index;第四遍看扩展:attention 各后端、engine/graph、moe、distributed/layers 的 TP 切分。配合 Nsight + NVTX 与 tests/ 下的单测,效果最好。

适合作为什么底座

官方定位已经说明:它是教学材料,也是研究原型。若要验证新的调度策略、新的 attention/MoE kernel、新的缓存层级或采样算法,mini-sglang 能让你在「与 SOTA 公平可比」的前提下只改局部;但任何要走向生产的功能(容错、跨机、长稳、多模态),最终还是应回到完整 SGLang。

总结一句话

mini-sglang 用约 5,000 行证明了一件事:现代 LLM 推理引擎的性能护城河,是 radix 前缀复用、分页内存、分块 prefill、重叠调度、CUDA Graph 与融合 kernel 这几近可数的机制;它们各自不过数百行代码,却共同撑起了与生产系统持平的性能。这正是最好的系统教材——复杂,但并非不可理解。

13

参考文献

References
  1. Mini-SGLang: Efficient Inference Engine in a NutshellZiyi Xu, LMSYS Org · 官方博客 · 2025 · 项目动机、架构图、benchmark 第一手来源lmsys.org
  2. mini-sglang GitHub 仓库sgl-project/mini-sglang · 本文所有源码引用的基线(python/minisgl)github
  3. SGLang: Efficient Execution of Structured Language Model ProgramsZheng et al. · NeurIPS 2024 · RadixAttention 与结构化语言程序arXiv:2312.07104
  4. SGLang: Fast and Expressive LLM Serving(RadixAttention 原始博客)LMSYS · 2024-01-17 · RadixAttention 概念与示例lmsys.org
  5. SGLang v0.4: Zero-Overhead Scheduler, Radix Cache, and MoreLMSYS · 2024-12-04 · overlap scheduling 工程实现与 Nsight 剖析lmsys.org
  6. Sarathi-Serve: Taming Throughput-Latency Tradeoff in LLM Inference ServingAgrawal et al. · OSDI 2024 · Chunked Prefill 与 stall-free batcharXiv:2403.02310
  7. NanoFlow: Towards Optimal Large Language Model Serving ThroughputHu et al. · 2024 · CPU/GPU 重叠调度流水线arXiv:2408.12757
  8. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessDao et al. · NeurIPS 2022 · tiling + online softmax,HBM 复杂度降到线性arXiv:2205.14135
  9. FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningDao · 2023 · 改进 warp 并行与工作划分arXiv:2307.08691
  10. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisionShah et al. · 2024 · Hopper WGMMA/TMA、异步、PDLarXiv:2407.08608
  11. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference ServingYe et al. · MLSys 2025 · 分页 KV、负载均衡 decode、CUDA Graph 兼容arXiv:2501.01005
  12. Efficient Memory Management for Large Language Model Serving with PagedAttention(vLLM)Kwon et al. · SOSP 2023 · 分页 KV 内存管理arXiv:2309.06180
  13. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head CheckpointsAinslie et al. · EMNLP 2023 · 分组 KV 头arXiv:2305.13245
  14. Megatron-LM: Training Multi-Billion Parameter Language Models Using Tensor ParallelismShoeybi et al. · 2019 · 张量并行切分与通信范式arXiv:1909.08053
  15. Nano-vLLMKedreamix · GitHub · 同类教学型迷你推理引擎,官方 benchmark 基线github
  16. TVM FFI(Foreign Function Interface)tlc-pack/tvm-ffi · mini-sglang 自研算子的轻量 Python 绑定与 JIT 基础设施github
滚轮缩放 · 拖拽平移 · 双击复位 · ESC 退出