Mini-SGLang

前四篇已经解释请求如何到达 Scheduler、KV Cache 如何管理，以及 Engine 怎样执行一个 Batch。最后还剩一个问题：几十亿参数怎样分布到多张 GPU，并在每层计算后重新组合？ 本文先推导 Tensor Parallel 的基本原理，再分析 Mini-SGLang 的 Layers、Models、Distributed、MoE 和 Kernel 模块。源码基于提交 9a91cfa。 一、为什么需要 Tensor Parallel 如果模型权重或 KV Cache 放不进单张 GPU，最直接的方法是把不同层放到不同 GPU，也就是 Pipeline Parallel。但自回归 Decode 每轮 token 很少，流水线容易出现气泡，而且跨 stage 传递激活会增加延迟。 Tensor Parallel 把同一层的矩阵切到多张 GPU，让所有 GPU 同时计算同一批 token： 同一 Transformer Layer GPU 0：处理部分权重 GPU 1：处理部分权重 GPU 2：处理部分权重 GPU 3：处理部分权重 ↓ collective 组合为完整层输出 它需要更频繁的 GPU 间通信，但能共同承载单层权重，并保持每个 rank 的执行进度一致。 二、Column Parallel Linear 考虑线性层： Y = XW 按输出维切分权重： ...

Scheduler 决定本轮执行哪些请求，Engine 则负责把这个决定变成 GPU 计算。它连接模型、KV Cache、Attention backend、CUDA Graph、Tensor Parallel 和采样器，是 Mini-SGLang 的执行核心。 本文先解释 Attention backend 与 CUDA Graph 的基本原理，再分析一次 Engine.forward_batch()。源码基于提交 9a91cfa。 一、GPU 推理不只有矩阵乘法 一次 Transformer forward 包含许多 GPU 操作： Embedding -> RMSNorm -> QKV Linear -> RoPE -> 写入 KV Cache -> Attention -> Output Linear -> Residual -> MLP / MoE -> LM Head -> Sampling 大模型 Prefill 中矩阵较大，GPU 容易被计算填满；Decode 中每个请求通常只有一个 query token，大量小 kernel 的 CPU launch 开销会变得突出。 ...

KV Cache 是在线 LLM 推理中最重要的状态。它避免 Decode 每轮重新计算全部历史 token，却也通常成为显存容量和并发数的主要限制。 本文先解释 KV Cache、分页管理和前缀复用的基本原理，再阅读 Mini-SGLang 的 MHAKVCache、CacheManager 与 RadixPrefixCache。源码基于提交 9a91cfa。 一、为什么需要 KV Cache 自回归 Attention 在第 t 步需要当前 Query 与位置 0...t 的 Key、Value 做注意力。如果每轮都重新计算历史 K、V，生成 n 个 token 会重复执行大量投影计算。 KV Cache 将每层历史 K、V 保存下来： 第 1 轮：计算 K0,V0，保存 第 2 轮：只计算 K1,V1，读取 [K0,K1] 第 3 轮：只计算 K2,V2，读取 [K0,K1,K2] 缓存大小大致为： 2 × 层数 × token 数 × KV head 数 × head_dim × 元素字节数 其中 2 表示 Key 和 Value。长上下文、大 batch 和高精度都会快速放大显存消耗。 ...

在线 LLM 调度器要回答的不是“下一个运行哪个进程”，而是：这一轮把哪些请求、哪些 token 放入同一个 GPU batch，并确保显存、KV Cache 和计算预算都不超限。 本文先建立 Continuous Batching、Chunked Prefill 与 Overlap Scheduling 的理论模型，再阅读 Mini-SGLang 的 Scheduler 实现。源码基于提交 9a91cfa。 一、调度器面对三种资源 LLM 请求至少消耗三类资源。 第一是计算量。Prefill 的计算大致随输入 token 数增长，Decode 每个请求每轮只增加一个 token。 第二是 KV Cache。一个请求即使本轮只计算一个 token，也必须继续持有全部历史 K、V。 第三是 batch 槽位和 CUDA Graph 形状。请求数、总 token 数和序列长度都会限制可执行 batch。 因此 Scheduler 需要同时维护两个预算： 本轮计算预算：最多处理多少新 token 长期缓存预算：这些请求最多还会占多少 KV 页 只看当前空闲页是不够的。如果 Prefill 阶段把缓存全部吃完，已经进入 Decode 的请求可能无法继续生成。 二、Continuous Batching 的状态划分 Mini-SGLang 将请求分成两个主要集合： waiting queue：尚未完成 Prefill，等待进入 GPU running batch：已经 Prefill，正在逐 token Decode 传统静态 batch 的生命周期以“整批”为单位；Continuous Batching 的生命周期以“请求”为单位。某个请求完成后立即离开，空出的槽位可以接纳新请求。 ...

大模型推理框架并不只是执行一次 model.forward()。在线服务面对的是持续到达、长度不同、生成进度不同的请求，它必须同时解决文本编解码、动态批处理、KV Cache、GPU 执行和流式返回。 Mini-SGLang 把这些问题压缩在一套相对紧凑的代码中。本文先不进入具体优化，而是建立阅读后续源码所需的整体模型：一次请求如何从 HTTP 文本进入系统，经过 Prefill 和多轮 Decode，最后以流式文本返回。 本文基于 Mini-SGLang 提交 9a91cfa。 一、推理为何分成 Prefill 和 Decode 输入 prompt 包含多个 token。模型第一次执行时，需要同时处理全部输入 token，并为每一层生成 Key、Value。这一阶段称为 Prefill。 输入： [t0, t1, t2, t3] 计算： 同时处理多个位置 产物： 最后位置的 logits + 四个位置的 KV Cache 接下来每轮只生成一个 token。已有 token 的 K、V 不应重复计算，只需读取缓存并计算新 token： 已有： [t0, t1, t2, t3] 第 1 轮 Decode：输入 t3，生成 t4 第 2 轮 Decode：输入 t4，生成 t5 第 3 轮 Decode：输入 t5，生成 t6 因此两阶段的计算形态不同：Prefill 计算量大、输入长度不等；Decode 单次计算小，但会循环很多轮。现代推理引擎通常分别为两者设计调度与 Attention kernel。 ...