Motivation

训练一个大模型很贵,但服务一个大模型可能更贵。

训练是一次性开销——几千张 GPU 跑几周就结束了。但推理是 7×24 小时不间断的:每一次用户输入、每一个 API 调用,都需要模型实时响应。OpenAI 每天处理数十亿个 token,推理成本占运营总成本的大头。如果推理效率翻倍,意味着成本直接减半——或者同样预算下服务两倍的用户。

推理和训练的优化方向截然不同。训练追求总计算吞吐——尽快把所有数据跑完;推理则需要同时优化两个相互矛盾的指标:

  • 吞吐量(Throughput):每秒生成多少个 token?决定了能服务多少并发用户
  • 延迟(Latency):每个用户等多久?TTFT(Time To First Token,首 token 延迟)和 TPOT(Time Per Output Token,逐 token 延迟)是关键指标

更重要的是,推理的计算特性和训练完全不同。我们在第一篇中分析过:自回归解码(autoregressive decoding)本质上是矩阵-向量乘法,算术强度只有约 1 FLOPs/Byte,GPU 算力利用率不到 1%。推理不是算力瓶颈,是显存带宽瓶颈——而 KV Cache 是其中最大的那块显存。

本文以 SGLang 为例,深入剖析推理系统的四个核心技术:

  1. KV Cache——推理显存的大头,为什么它会成为瓶颈
  2. PagedAttention——借鉴 OS 虚拟内存思想,解决 KV Cache 的显存碎片问题
  3. RadixAttention——SGLang 的核心创新,用 Radix Tree 实现 KV Cache 复用
  4. Continuous Batching——迭代级调度,最大化 GPU 利用率
  5. Chunked Prefill——解决长 prompt 阻塞解码的问题

读完之后,你应该能理解为什么 vLLM 和 SGLang 能比朴素实现快 2-10 倍,以及它们之间的设计差异。

前置知识

  • GPU 显存模型与分布式通信基础(第 1 篇)——特别是 HBM 带宽瓶颈和 Roofline 模型
  • Transformer 注意力机制的基本理解(Self-Attention、Q/K/V 投影)
  • 了解自回归生成的基本流程(逐 token 生成)

KV Cache 基础

为什么需要 KV Cache

自回归语言模型在生成第 $t$ 个 token 时,Attention 的计算公式是:

$$\text{Attention}(Q_t, K_{1:t}, V_{1:t}) = \text{softmax}\left(\frac{Q_t \cdot K_{1:t}^T}{\sqrt{d_k}}\right) V_{1:t}$$

关键点:当前 token 的 Query 需要和所有已生成 token 的 Key、Value 做内积。如果每次都从头计算所有 token 的 K 和 V,复杂度是 $O(t^2)$——生成 1000 个 token 就要做 $1 + 2 + \cdots + 1000 = 500500$ 次 attention 计算。

KV Cache 的核心思想:之前 token 的 K 和 V 不会变,只需要计算一次,缓存起来复用。

flowchart TD
    subgraph NO["自回归解码(无 KV Cache)"]
        N1["Step 1: 输入 [t1] → 计算 K1,V1 → Attention → 输出 t2"]
        N2["Step 2: 输入 [t1,t2] → 重新计算 K1,V1,K2,V2 → Attention"]
        N3["Step 3: 输入 [t1,t2,t3] → 重新计算 K1,K2,K3... → Attention"]
        NN["Step N: 输入 [t1,...,tN] → 重新计算所有 K,V → O(N²) 总计算"]
        N1 --> N2 --> N3 -.-> NN
    end

    subgraph YES["自回归解码(有 KV Cache)"]
        Y1["Step 1: 输入 [t1] → 计算 K1,V1 → 缓存 → Attention → 输出 t2"]
        Y2["Step 2: 输入 [t2] → 计算 K2,V2 → 追加缓存 → Attention"]
        Y3["Step 3: 输入 [t3] → 计算 K3,V3 → 追加缓存 → Attention"]
        YN["Step N: 输入 [tN] → 计算 KN,VN → 追加缓存 → O(N) 总计算"]
        Y1 --> Y2 --> Y3 -.-> YN
        YNOTE["每步只计算 1 个 token 的 K,V,复用缓存中的所有 K,V"]
    end

    NO --> YES

    style NO fill:#fff3cd,stroke:#856404
    style YES fill:#d4edda,stroke:#155724

KV Cache 将总计算量从 $O(N^2)$ 降到了 $O(N)$——这是推理优化的第一性原理,所有现代推理引擎都必须实现它。

KV Cache 有多大

每个 token 的 KV Cache 大小为:

$$\text{KV per token} = 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{dtype_size}$$

其中因子 2 是因为要存 K 和 V 两份。以几个主流模型为例(FP16,2 bytes per element):

模型$n_{\text{layers}}$$n_{\text{heads}}$$d_{\text{head}}$KV per token2048 tokens
LLaMA-7B32321280.5 MB1.0 GB
LLaMA-13B40401280.8 MB1.6 GB
LLaMA-70B80641282.5 MB5.0 GB
GPT-3 175B96961284.5 MB9.2 GB

一个 LLaMA-70B 的请求如果生成 2048 个 token,仅 KV Cache 就占 5 GB。如果你想同时服务 16 个这样的请求,KV Cache 就需要 80 GB——一整张 A100 的显存。而且模型权重本身还要 140 GB(FP16)

这就是推理显存的核心矛盾:模型权重是固定开销,KV Cache 是动态开销,而 KV Cache 随并发请求数和序列长度线性增长,很容易成为瓶颈。

能同时 batch 多少请求,取决于 KV Cache 能放多少——batch size 直接决定吞吐量。所以 KV Cache 的内存管理效率,就是推理系统的性能上限。

朴素分配的浪费

最直观的做法是:为每个请求预分配 max_seq_len 大小的连续显存来存 KV Cache。问题是,你不知道一个请求到底会生成多少 token——可能是 10 个,也可能是 2000 个。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class NaiveKVCache:
    """
    Pre-allocated KV cache for a single sequence.

    Problem: if max_seq_len=2048 but the sequence only generates 50 tokens,
    we waste 2048-50 = 1998 slots of memory.  Multiply by batch_size and
    num_layers, and you're throwing away most of your GPU memory.
    """

    def __init__(self, max_seq_len: int, num_heads: int, head_dim: int):
        self.max_seq_len = max_seq_len
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.current_len = 0

        # Pre-allocate full tensors — this is the waste!
        self.k_cache = torch.zeros(max_seq_len, num_heads, head_dim, device=device)
        self.v_cache = torch.zeros(max_seq_len, num_heads, head_dim, device=device)

配套代码中的实验表明,当请求长度分布不均匀时(现实中几乎总是如此),朴素预分配的显存浪费率高达 60-80%

5 requests, max_seq_len=2048
Actual lengths: [37, 152, 8, 420, 91]
Memory allocated: 160.0 MB
Memory actually used:  22.1 MB
Waste: 86.2%

--> This is why we need PagedAttention!

显存浪费 86%——这意味着原本能同时处理 7 个请求的 GPU,实际上只能处理 1 个。浪费的不仅是内存,更是吞吐量和收入。

KV Cache 增长与浪费


PagedAttention

从操作系统借鉴的智慧

PagedAttention 是 vLLM(OSDI'23)的核心贡献,它的灵感来自一个成熟了 50 年的技术——操作系统的虚拟内存

在操作系统中,进程看到的是虚拟地址空间——连续的、独立的。但物理内存是以固定大小的页(page) 分配的,通过页表(page table) 将虚拟地址映射到物理地址。进程不需要拿到连续的物理内存,操作系统在幕后把碎片化的物理页拼成看似连续的虚拟空间。

PagedAttention 把这个思想完美移植到 KV Cache 管理上:

flowchart LR
    subgraph OS["OS 虚拟内存"]
        A1["进程"]
        A2["虚拟页"]
        A3["物理页帧"]
        A4["页表"]
        A5["按需分配(缺页中断)"]
        A6["页大小固定(4KB)"]
        A7["CoW(fork)"]
        A8["页面置换(swap)"]
    end
    subgraph PA["PagedAttention"]
        B1["请求(sequence)"]
        B2["逻辑 KV block"]
        B3["物理 KV block"]
        B4["block table"]
        B5["按需分配(生成新 token 时)"]
        B6["block 大小固定(16 tokens)"]
        B7["CoW(beam search)"]
        B8["preemption(换出请求)"]
    end
    A1 -.-> B1
    A2 -.-> B2
    A3 -.-> B3
    A4 -.-> B4
    A5 -.-> B5
    A6 -.-> B6
    A7 -.-> B7
    A8 -.-> B8

    style OS fill:#cce5ff,stroke:#004085
    style PA fill:#d4edda,stroke:#155724

Block 分配机制

PagedAttention 将 GPU 上的 KV Cache 显存划分为一个固定大小的 block 池。每个 block 能存储固定数量的 token 的 KV(通常 16 或 32 个 token)。每个请求有一个 block table,记录其逻辑 block 到物理 block 的映射。

flowchart TD
    subgraph POOL["物理 Block 池(GPU 显存)"]
        B0["B0<br/>Seq0"]
        B1["B1<br/>Seq1"]
        B2["B2<br/>Free"]
        B3["B3<br/>Seq0"]
        B4["B4<br/>Seq1"]
        B5["B5<br/>Free"]
        B6["B6<br/>Seq0"]
        B7["B7<br/>Free"]
    end

    subgraph T0["Seq 0 Block Table"]
        L0_0["Logic 0 → B0"]
        L0_1["Logic 1 → B3"]
        L0_2["Logic 2 → B6"]
    end

    subgraph T1["Seq 1 Block Table"]
        L1_0["Logic 0 → B1"]
        L1_1["Logic 1 → B4"]
    end

    L0_0 --> B0
    L0_1 --> B3
    L0_2 --> B6
    L1_0 --> B1
    L1_1 --> B4

    NOTE["物理 block 不需要连续,逻辑上连续即可"]

    style POOL fill:#cce5ff,stroke:#004085
    style B2 fill:#f8f9fa,stroke:#6c757d
    style B5 fill:#f8f9fa,stroke:#6c757d
    style B7 fill:#f8f9fa,stroke:#6c757d

工作流程:

  1. 新请求到达:分配一个空的 block table,不预分配任何 block
  2. 生成第 1 个 token:从 free pool 拿一个 block,存入 KV,更新 block table
  3. 当前 block 填满(写满 16 个 token 的 KV):再从 free pool 拿一个新 block
  4. 请求结束:归还所有 block 到 free pool,立即可被其他请求复用
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class BlockManager:
    """
    PagedAttention-style block manager.

    Key ideas:
      - Physical KV cache is a pool of fixed-size blocks.
      - Each sequence has a block table mapping logical -> physical blocks.
      - Blocks are allocated on demand as sequences grow.
      - Free blocks are recycled when sequences finish.
      - Copy-on-Write: shared blocks are copied only when modified.
    """

    def __init__(self, num_blocks, num_layers, num_heads, head_dim, block_size=16):
        self.num_blocks = num_blocks
        self.block_size = block_size

        # Physical block pool — the actual GPU memory
        self.k_pool = torch.zeros(
            num_layers, num_blocks, block_size, num_heads, head_dim, device=device
        )
        self.v_pool = torch.zeros(
            num_layers, num_blocks, block_size, num_heads, head_dim, device=device
        )

        # Track physical block metadata
        self.physical_blocks = [PhysicalBlock(block_id=i) for i in range(num_blocks)]
        self.free_block_ids: list[int] = list(range(num_blocks))
        self.seq_tables: dict[int, SequenceBlockTable] = {}

从 86% 浪费到 4% 浪费

PagedAttention 的内存浪费只来自一个地方:每个请求最后一个 block 的内部碎片(internal fragmentation)。如果 block_size=16,最后一个 block 平均只用了一半,浪费约 $\frac{1}{2} \times \frac{1}{N_{\text{blocks}}}$ 的显存。

配套代码中的对比实验结果:

  Config: 32 heads, head_dim=128, max_seq=2048, block_size=16
  10 requests with lengths: [47, 183, 12, 891, 256, 5, 1024, 73, 330, 15]

  Metric                              Naive        Paged       Actual
  -----------------------------------------------------------------------
  KV cache memory (MB)               320.0         45.5         44.3
  Memory utilization                  13.8%        97.4%       100.0%
  Waste (MB)                         275.7          1.2          0.0
  Waste (%)                          86.2%         2.6%         0.0%

朴素方式浪费 86%,PagedAttention 只浪费 2.6%。这意味着同样的 GPU 显存可以 batch 多 2-4 倍的请求,直接把吞吐量提高 2-4 倍。

PagedAttention block 管理

Copy-on-Write:高效的序列分叉

在 beam search 和 parallel sampling 场景下,一个请求会分叉出多个候选序列。这些序列共享前缀的 KV Cache——如果每个候选都复制一份前缀的 KV,显存开销会翻倍。

PagedAttention 的解决方案是 Copy-on-Write(CoW),同样借鉴自操作系统的 fork() 系统调用:

flowchart TD
    subgraph BEFORE["fork 前"]
        S0a["Seq 0: Block A → Block B → Block C (ref=1)"]
    end

    subgraph AFTER["fork 后(Seq 0 + Seq 1)"]
        S0b["Seq 0: Block A → Block B → Block C"]
        S1b["Seq 1: Block A → Block B → Block C"]
        REF["Block C ref=2(共享,无拷贝!)"]
    end

    subgraph WRITE["Seq 1 写入新 token 到 Block C"]
        S0c["Seq 0: Block A → Block B → Block C (ref=1)"]
        S1c["Seq 1: Block A → Block B → Block C' (ref=1, 此时拷贝)"]
        NOTE2["只有被修改的 block 才拷贝,前面的 block 仍共享"]
    end

    BEFORE --> AFTER --> WRITE

    style BEFORE fill:#cce5ff,stroke:#004085
    style AFTER fill:#fff3cd,stroke:#856404
    style WRITE fill:#d4edda,stroke:#155724

核心逻辑在 append_token 中:写入前检查 block 的引用计数,如果 >1(被共享),就先拷贝一份新 block,再写入:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def append_token(self, seq_id, k, v, layer_idx):
    table = self.seq_tables[seq_id]
    logical_block_idx = token_idx // self.block_size
    slot_in_block = token_idx % self.block_size

    # Allocate new block if current one is full
    if logical_block_idx >= len(table.logical_to_physical):
        phys_id = self._allocate_block()
        table.logical_to_physical.append(phys_id)

    phys_id = table.logical_to_physical[logical_block_idx]

    # CoW check: if this block is shared, copy before writing
    if self.physical_blocks[phys_id].ref_count > 1:
        new_phys_id = self._allocate_block()
        self.k_pool[:, new_phys_id] = self.k_pool[:, phys_id]
        self.v_pool[:, new_phys_id] = self.v_pool[:, phys_id]
        self._free_block(phys_id)  # Decrement old block's ref count
        table.logical_to_physical[logical_block_idx] = new_phys_id
        phys_id = new_phys_id

    # Write KV to the physical block
    self.k_pool[layer_idx, phys_id, slot_in_block] = k
    self.v_pool[layer_idx, phys_id, slot_in_block] = v

CoW 的收益在 beam search 中尤其显著:beam_size=4 时,4 个候选序列共享大量前缀 KV,只有最后一个 block(正在生成的部分)需要独立存储。显存开销接近 1 个序列而非 4 个。


RadixAttention(SGLang 核心创新)

从 PagedAttention 到 Prefix Caching

PagedAttention 解决了单个请求内部的显存碎片问题。但它忽略了一个更大的优化机会:不同请求之间的 KV Cache 复用。

在实际的 LLM 应用中,大量请求共享相同的前缀

flowchart TD
    subgraph S1["场景 1: System Prompt"]
        SP["共享前缀: You are a helpful assistant..."]
        UA["+ 用户 A: What is machine learning?"]
        UB["+ 用户 B: Write a Python function for..."]
        UC["+ 用户 C: Translate this text..."]
        SP --> UA
        SP --> UB
        SP --> UC
    end

    subgraph S2["场景 2: Few-shot Learning"]
        FS["共享前缀: Here are some examples: × 5 样例"]
        Q1["+ Input: new query 1"]
        Q2["+ Input: new query 2"]
        FS --> Q1
        FS --> Q2
    end

    subgraph S3["场景 3: Multi-turn Chat"]
        MT["共享前缀: System + Turn 1 + Turn 2"]
        T3["+ User Turn 3(新消息)"]
        MT --> T3
    end

    style S1 fill:#cce5ff,stroke:#004085
    style S2 fill:#d4edda,stroke:#155724
    style S3 fill:#fff3cd,stroke:#856404

如果 system prompt 有 500 个 token,每个请求都重新计算和存储这 500 个 token 的 KV Cache,那么 100 个并发请求就重复存了 100 份——完全相同的数据。

Prefix Caching 的核心思想:如果两个请求的 token 序列有相同的前缀,它们的 KV Cache 在这些位置上完全一样,可以直接共享,无需重新计算。

vLLM 后来也加入了 prefix caching 功能,但它的实现基于预定义的前缀匹配——你需要显式声明哪些请求共享前缀。SGLang 的 RadixAttention 提供了一种更灵活、更优雅的方案。

Radix Tree:索引任意前缀的数据结构

SGLang(LMSYS,2024)的核心创新是用 Radix Tree(基数树) 来索引和管理所有已缓存的 KV Cache。

Radix Tree 是一种压缩前缀树(compressed trie)。普通 trie 每个节点代表一个字符,radix tree 将只有一个子节点的路径压缩成单个节点,减少树的深度。在 SGLang 中,每个节点代表一段 token 序列,节点上存储着对应的 KV Cache block 指针。

graph TD
    ROOT["Root"]
    SYS["System Prompt<br/>KV blocks: B0, B1, B2"]
    ML["User: What is ML?<br/>KV blocks: B3, B4"]
    PY["User: Write Python<br/>KV blocks: B5, B6"]
    AST["Asst: ML is...<br/>KV blocks: B7"]

    ROOT --> SYS
    SYS --> ML
    SYS --> PY
    ML --> AST

    REQD["新请求 D: System Prompt + What is ML? + new query<br/>→ 匹配 B0-B4 直接复用,只需为 new query 计算新 KV"]

    style SYS fill:#cce5ff,stroke:#004085
    style ML fill:#d4edda,stroke:#155724
    style PY fill:#d4edda,stroke:#155724
    style AST fill:#fff3cd,stroke:#856404
    style REQD fill:#f8f9fa,stroke:#6c757d

Radix Tree 构建过程

工作流程

当一个新请求到达时,SGLang 的调度器执行以下步骤:

flowchart TD
    subgraph FLOW["RadixAttention 处理新请求"]
        A["1. PREFIX MATCHING<br/>在 Radix Tree 中查找最长匹配前缀<br/>例: sys_prompt 的 KV blocks 可直接复用"]
        B["2. REUSE CACHED KV<br/>匹配到的前缀部分跳过 prefill<br/>只需计算不匹配的后缀部分的 KV"]
        C["3. INSERT INTO TREE<br/>请求完成后将新 KV 插入 Radix Tree<br/>后续相同前缀的请求可复用"]
        A --> B --> C
    end

    subgraph TIMELINE["时间对比"]
        T1["无 cache: ====== prefill 500 tokens ======  decode"]
        T2["有 cache: == prefill 100 tokens ==  decode<br/>400 tokens 的 KV 从 cache 复用,跳过计算"]
    end

    FLOW --> TIMELINE

    style FLOW fill:#cce5ff,stroke:#004085
    style T1 fill:#fff3cd,stroke:#856404
    style T2 fill:#d4edda,stroke:#155724

在 system prompt 场景下(几乎所有请求共享同一个 system prompt),prefix caching 可以跳过大部分 prefill 计算,将 TTFT(Time To First Token)降低 50-90%

LRU 逐出策略

GPU 显存有限,不可能缓存所有见过的前缀。当 block 池用满时,需要逐出一些缓存来给新请求腾空间。

SGLang 使用 LRU(Least Recently Used) 策略:最久没被访问的 KV Cache 节点优先被逐出。

flowchart TD
    subgraph NODES["Radix Tree 节点(按最近访问时间排序)"]
        N1["sys_prompt + user_A — 10 sec ago ✅ 保留"]
        N2["sys_prompt + user_B — 30 sec ago ✅ 保留"]
        N3["sys_prompt + user_C — 120 sec ago ⚠️ 候选逐出"]
        N4["old_prompt + old_query — 300 sec ago ❌ 优先逐出"]
    end

    subgraph EVICT["逐出流程"]
        E1["1. 从 LRU 链表尾部取最久未访问的节点"]
        E2["2. 释放该节点的 KV blocks"]
        E3["3. 从 Radix Tree 中删除该节点"]
        E4["4. 若父节点变成无引用叶子,递归清理"]
        E1 --> E2 --> E3 --> E4
    end

    style N1 fill:#d4edda,stroke:#155724
    style N2 fill:#d4edda,stroke:#155724
    style N3 fill:#fff3cd,stroke:#856404
    style N4 fill:#f8d7da,stroke:#721c24

LRU 的合理性在于时间局部性:最近被访问的前缀往往会再次被请求到(比如同一个用户的多轮对话,或热门的 system prompt)。

Cache-Aware Scheduling

有了 Radix Tree 之后,调度器可以做一个非常聪明的优化:优先调度那些在 cache 中匹配前缀更长的请求

flowchart TD
    subgraph QUEUE["等待队列"]
        RA["Req A: 匹配 500 tokens 缓存 → prefill 仅需 100 tokens"]
        RB["Req B: 匹配 0 tokens → prefill 需要 600 tokens"]
        RC["Req C: 匹配 300 tokens 缓存 → prefill 仅需 200 tokens"]
    end

    subgraph SCHED["调度策略对比"]
        FCFS["朴素 FCFS: A → B → C"]
        CACHE["Cache-aware: A → C → B(优先 cache 命中率高的)"]
    end

    subgraph BENEFIT["好处"]
        B1["1. 高 cache 命中的请求 prefill 更快 → TTFT 更低"]
        B2["2. KV 留在 tree 中 → 后续相似请求也命中 → 正反馈循环"]
        B3["3. 减少总 prefill 计算量 → 更多 GPU 时间给 decode"]
    end

    QUEUE --> SCHED --> BENEFIT

    style RA fill:#d4edda,stroke:#155724
    style RB fill:#f8d7da,stroke:#721c24
    style RC fill:#fff3cd,stroke:#856404
    style CACHE fill:#d4edda,stroke:#155724

这种调度策略使 SGLang 在 few-shot learning、multi-turn chat 等共享前缀场景下性能远超不做 prefix caching 的系统。

与 vLLM 的对比

vLLM(v0.3+)后来也加入了 prefix caching 功能,但设计思路不同:

特性SGLang (RadixAttention)vLLM (Prefix Caching)
数据结构Radix TreeHash Map
前缀匹配自动匹配任意共享前缀基于 token block hash
匹配粒度token 级别block 级别(block_size 对齐)
逐出策略LRU on tree nodesLRU on blocks
Cache-aware 调度原生支持后续版本加入
多轮对话天然支持(树结构)需要 hash 匹配

SGLang 的 Radix Tree 方案在灵活性上更胜一筹:

  • 任意前缀共享:不限于预定义的 prompt template,任何两个请求只要有共同前缀就能自动匹配
  • 树结构天然适配多轮对话:对话的历史消息自然形成一棵树,不同分支代表不同的后续对话
  • 渐进式匹配:可以匹配到 token 级别的精确前缀,而非 block 级别的近似匹配

vLLM 的 hash-based 方案在实现简单性大规模部署稳定性上有优势,两者在实际性能上差距取决于具体的工作负载特征。


Continuous Batching

静态 Batching 的问题

在理解了 KV Cache 的内存管理之后,我们来看另一个维度的优化:调度策略

最朴素的推理方式是 静态 batching:把一组请求凑成一个 batch,一起做 prefill,然后一起做 decode,等 batch 中所有请求都生成完毕后,才处理下一个 batch。

问题在于:不同请求的输出长度差异巨大。

flowchart TD
    subgraph STATIC["静态 Batching — Batch = Req 0-3, 输出长度: 3, 8, 2, 5"]
        direction LR
        R0["Req 0: ## ## ## .. .. .. .. ..<br/>第 3 步结束"]
        R1["Req 1: ## ## ## ## ## ## ## ##<br/>最长,其他人都在等"]
        R2["Req 2: ## ## .. .. .. .. .. ..<br/>第 2 步结束"]
        R3["Req 3: ## ## ## ## ## .. .. ..<br/>第 5 步结束"]
    end

    WASTE["## = 有效计算 · .. = GPU 空闲<br/>GPU 利用率: 18/32 = 56.25%<br/>浪费了 43.75% 的 GPU 算力<br/>Req 4, 5 必须等 Req 1 结束才能开始"]

    style STATIC fill:#fff3cd,stroke:#856404
    style WASTE fill:#f8d7da,stroke:#721c24

当输出长度方差很大时(这在实际中几乎总是如此——有人问"是几点了",有人要写一篇长文),静态 batching 的 GPU 利用率可以低至 20-50%。

迭代级调度

Continuous Batching(Orca, OSDI'22)的核心改进是将调度粒度从 batch 级别 降到 iteration 级别

每一个 decode step 结束后,立即检查:

  1. 哪些请求已经结束(生成了 EOS 或达到 max_length)?→ 驱逐,释放 KV Cache
  2. 等待队列中有没有新请求可以加入?→ 纳入,开始 prefill
flowchart TD
    subgraph CB["Continuous Batching — 迭代级调度"]
        direction LR
        R0["Req 0: ## ## ## — step 3 结束"]
        R1["Req 1: ## ## ## ## ## ## ## ## — step 8 结束"]
        R2["Req 2: ## ## — step 2 结束"]
        R3["Req 3: ## ## ## ## ## — step 5 结束"]
        R4["Req 4: step 3 加入 → ## ## ## ##"]
        R5["Req 5: step 4 加入 → ## ## ## ## ## ##"]
    end

    RESULT["每步活跃数: 4 4 4 4 4 3 3 3 1<br/>GPU 几乎一直满载!结束的请求立刻被新请求替换"]

    style CB fill:#d4edda,stroke:#155724
    style RESULT fill:#cce5ff,stroke:#004085

配套代码中的对比实验显示,continuous batching 在吞吐量上通常有 2-3 倍的提升:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def run_continuous_batching(requests, max_batch_size, total_blocks=256, ...):
    """
    After EVERY decode step:
      1. Remove finished sequences (free their KV blocks).
      2. If there are free slots AND enough KV blocks, admit waiting requests.
      3. If out of KV blocks, preempt the newest running sequence (LIFO).
    """
    # ... 每步检查并调度 ...
    for req in running:
        req.tokens_generated += 1
        if req.tokens_generated >= req.target_output_length:
            req.status = SequenceStatus.FINISHED
            finished_this_step.append(req)

    # Remove finished, free their blocks
    for req in finished_this_step:
        running.remove(req)
        block_pool.free(req.num_kv_blocks)

    # Try to admit waiting requests
    while waiting and len(running) < max_batch_size:
        req = waiting[0]
        if block_pool.can_allocate(blocks_needed):
            waiting.popleft()
            running.append(req)

请求调度流程

Preemption:优雅应对显存压力

Continuous batching 还引入了一个关键机制:preemption(抢占)。当 KV Cache 显存不够为所有 running 请求分配新 block 时,调度器可以:

  1. 暂停一个或多个 running 请求(通常选最新加入的——LIFO 策略,因为它们做的工作最少)
  2. 释放被暂停请求的 KV Cache blocks
  3. 用腾出的空间继续服务其他请求
  4. 稍后恢复被暂停的请求(重新做 prefill 来重建 KV Cache)
flowchart TD
    A["正常运行<br/>Running: Req 0, 1, 2, 3<br/>KV blocks: 200/256 (78%)"]
    B["Req 1 生成长回复, KV 持续增长...<br/>KV blocks: 253/256 (99%) — 即将 OOM!"]
    C["触发 Preemption<br/>暂停 Req 3 (LIFO: 最新, 做的工作最少)<br/>释放 Req 3 的 20 个 blocks"]
    D["KV blocks: 233/256 (91%)<br/>有空间继续了"]
    E["稍后 Req 0/1 结束, 空间释放<br/>→ 恢复 Req 3, 重新 prefill"]
    F["替代方案: 不做 preemption → OOM crash → 所有请求失败!"]

    A --> B --> C --> D --> E
    B -.-> F

    style A fill:#d4edda,stroke:#155724
    style B fill:#f8d7da,stroke:#721c24
    style C fill:#fff3cd,stroke:#856404
    style D fill:#d4edda,stroke:#155724
    style F fill:#f8d7da,stroke:#721c24

Preemption 的代价是被暂停的请求需要重新做 prefill(重建 KV Cache),增加了该请求的延迟。但比起 OOM 导致所有请求崩溃,这是一个合理的 trade-off。

配套代码中的 preemption demo 显示,即使在非常受限的显存下(32 blocks),preemption 也能保证所有请求最终完成:

Config: 20 requests, 32 KV blocks, max_batch=4
Preemptions triggered: 12
Preemption allows the system to handle load spikes without OOM crashes.
All 20 requests completed: True

Chunked Prefill

长 Prompt 阻塞解码

Prefill 阶段(处理用户输入的 prompt)和 decode 阶段(逐 token 生成)有截然不同的计算特性:

  • Prefill:一次性处理整个 prompt(可能数千个 token),是 Compute-bound 的大 GEMM 操作
  • Decode:每步只生成 1 个 token,是 Memory-bound 的 GEMV 操作

问题在于:如果一个请求的 prompt 有 8000 个 token,做 prefill 可能需要几百毫秒。在这段时间里,所有正在 decode 的请求都被阻塞了——它们在等 GPU 做完这个长 prefill。

flowchart TD
    subgraph NOCHUNK["无 Chunked Prefill"]
        GPU1["GPU: ====== Prefill 8K tokens (200ms) ====== → D D D D"]
        RA1["Req A: 等待 200ms... 然后才能 decode"]
        RB1["Req B: 等待 200ms... 然后才能 decode"]
    end

    PROBLEM["问题: Req A, B 的 TPOT 从 10ms 暴增到 200ms+<br/>用户体验: 正在打字的对话突然卡住 200ms"]

    style NOCHUNK fill:#fff3cd,stroke:#856404
    style PROBLEM fill:#f8d7da,stroke:#721c24

对于正在和模型对话的用户来说,TPOT 从 10ms 跳到 200ms 会造成明显的"卡顿感"——字在流式输出时突然停了一下。

分块 Prefill

Chunked Prefill 的解决方案直截了当:把长 prompt 的 prefill 分成多个小块(chunk),每个 chunk 处理固定数量的 token(比如 512 个),每个 chunk 之间穿插一轮 decode step。

flowchart TD
    subgraph CHUNK["Chunked Prefill — 8192 tokens, chunk_size=512, 共 16 chunk"]
        GPU2["GPU: P1 → D → P2 → D → P3 → D → P4 → D → P5 ..."]
        LEGEND["P = prefill chunk (512 tokens, ~12ms) · D = decode step (~3ms)"]
    end

    subgraph DECODE["Decode 请求视角"]
        RA2["Req A: D — D — D — D — D (每 ~15ms 一次)"]
        RB2["Req B: D — D — D — D — D (每 ~15ms 一次)"]
    end

    RESULT2["TPOT: ~15ms(稳定!)vs 无 chunking 时的 200ms 尖峰"]

    CHUNK --> DECODE
    DECODE --> RESULT2

    style CHUNK fill:#d4edda,stroke:#155724
    style RESULT2 fill:#cce5ff,stroke:#004085

Trade-off

  • Prefill 总时间变长了:因为每个 chunk 之间插入了 decode step,8K token 的 prefill 从 200ms 变成了约 240ms(TTFT 略增)
  • Decode 延迟变稳定了:TPOT 从 200ms 的尖峰降到稳定的 15ms

这是一个 TTFT vs TPOT 的 trade-off:

指标无 Chunked Prefill有 Chunked Prefill
TTFT (首 token)200ms~240ms (+20%)
TPOT (逐 token,最坏情况)200ms (被阻塞)~15ms (稳定)
TPOT (P99)极差可控
用户体验偶尔严重卡顿流畅稳定

在大多数在线服务场景下,稳定的 TPOT 比略低的 TTFT 更重要——用户更在意"打字流畅"而非"快 40ms 看到第一个字"。

Chunk 大小的选择

Chunk 大小的选择需要平衡:

  • 太小(如 64 tokens):prefill 的 GPU 利用率低(小矩阵乘法效率差),TTFT 显著增加
  • 太大(如 4096 tokens):接近不分块的情况,decode 仍然会被长时间阻塞
  • 典型值:512-1024 tokens,兼顾 prefill 效率和 decode 稳定性

SGLang 和 vLLM 都支持 chunked prefill,并可根据当前 running batch 的状态动态调整 chunk 大小——如果当前没有 decode 请求,就不需要分块,直接做完整的 prefill 以最小化 TTFT。


推理系统全景

将以上所有组件整合起来,一个现代 LLM 推理引擎(如 SGLang)的整体架构如下:

flowchart TD
    HTTP["HTTP Server<br/>接收用户请求"]

    subgraph SCHED["Scheduler(调度器 — 大脑)"]
        WQ["Waiting Queue"]
        RT["Radix Tree<br/>(Prefix Cache 索引)"]
        STEPS["每个 iteration:<br/>1. 驱逐已完成请求<br/>2. 匹配新请求前缀<br/>3. Cache-aware 排序 + 纳入<br/>4. 必要时 preempt<br/>5. Chunked prefill 调度"]
    end

    subgraph BM["Block Manager(块管理器 — 内存)"]
        POOL2["Physical Block Pool: B0 B1 B2 B3 B4 B5 ..."]
        BMOPS["按需分配/释放 · CoW · LRU 逐出"]
    end

    subgraph ME["Model Executor(模型执行器 — 计算)"]
        ATTN["Attention kernel (FlashAttention / FlashInfer)"]
        KV["从 Block Table 读取 KV(非连续物理地址)"]
        PAR["TP / PP 并行推理"]
    end

    HTTP --> SCHED --> BM --> ME

    style SCHED fill:#cce5ff,stroke:#004085
    style BM fill:#fff3cd,stroke:#856404
    style ME fill:#d4edda,stroke:#155724

调度器是大脑——它决定每一步执行哪些请求,是 prefill 还是 decode,是否需要 preempt。Block Manager 是内存管理器——它负责分配、释放、共享 KV Cache blocks。Model Executor 是执行引擎——它拿到调度器的决策和 block table 的映射,执行实际的 attention 计算。

三者协同工作,构成了一个高效的推理流水线。


Key Takeaways

  1. KV Cache 是推理的显存瓶颈。每个 token 需要存储 $2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{dtype_size}$ 的 KV 数据。LLaMA-70B 下每个 token 约 2.5 MB,100 个并发请求、每个 2048 tokens 就需要 500 GB 的 KV Cache——远超单卡显存。能 batch 多少请求、能处理多长序列,完全取决于 KV Cache 的管理效率。

  2. PagedAttention 用 OS 虚拟内存思想解决显存碎片。固定大小的 block 按需分配,block table 做逻辑-物理映射,CoW 支持高效的序列分叉。将显存浪费从 60-80% 降到 <5%,同样的 GPU 可以 batch 2-4 倍的请求。

  3. RadixAttention 实现跨请求的 KV Cache 复用。Radix Tree 索引所有已缓存的 token 前缀,新请求自动匹配最长前缀并复用其 KV Cache。结合 LRU 逐出和 cache-aware scheduling,在共享前缀场景下(system prompt、few-shot、multi-turn chat)大幅降低 TTFT。

  4. Continuous Batching 将调度粒度从 batch 级降到 iteration 级。每步驱逐已完成请求、纳入等待请求,GPU 利用率从 20-50% 提升到 80-95%。配合 preemption 机制,优雅应对显存压力,避免 OOM 崩溃。

  5. Chunked Prefill 解决长 prompt 阻塞 decode 的问题。将 prefill 分成小块,穿插 decode step,用略高的 TTFT 换取稳定的 TPOT。chunk 大小的选择需要根据负载特征权衡。

  6. 这些技术不是互斥的,而是层层叠加的。现代推理引擎(SGLang、vLLM)同时使用 PagedAttention + Prefix Caching + Continuous Batching + Chunked Prefill,每一层解决不同维度的效率问题,叠加起来实现了比朴素实现 2-10 倍的吞吐量提升。


配套代码

本文配套代码位于 code/03-inference-sglang/

  • kv_cache_from_scratch.py — 从零实现 KV Cache 和 PagedAttention Block Manager。包含 4 个部分:(1) 朴素 KV Cache 的浪费演示;(2) PagedAttention block 分配、释放、CoW 的完整实现;(3) 用 block-managed KV cache 跑一个 mini decode loop;(4) 朴素 vs PagedAttention 的内存效率对比。运行方式:python kv_cache_from_scratch.py

  • continuous_batching.py — 静态 batching vs continuous batching 的对比模拟。包含 4 个部分:(1) 静态/连续 batching 的吞吐量对比;(2) ASCII 调度时间线可视化;(3) preemption 在显存压力下的表现;(4) 不同 batch size 下的 scaling 分析。运行方式:python continuous_batching.py

所有代码支持 CPU 模式(无需 GPU),方便在任何环境下学习核心逻辑。


参考资料

  1. Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP'23 (vLLM) — PagedAttention 的原始论文,将 OS 虚拟内存思想引入 KV Cache 管理
  2. SGLang: Efficient Execution of Structured Language Model Programs — Zheng et al., 2024 (LMSYS) — RadixAttention 和 cache-aware scheduling 的原始论文
  3. Orca: A Distributed Serving System for Transformer-Based Generative Models — Yu et al., OSDI'22 — Continuous Batching(iteration-level scheduling)的开创性工作
  4. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al., NeurIPS'22 — 推理引擎中 attention kernel 的核心优化
  5. vLLM Documentationdocs.vllm.ai — vLLM 官方文档
  6. SGLang Documentationsgl-project.github.io — SGLang 官方文档
  7. Sarathi-Serve: Chunked-Prefills for Efficient LLM Serving — Agrawal et al., 2024 — Chunked Prefill 的系统化分析