← 返回目录 模块 2 / 8

2.1 PagedAttention

核心创新

PagedAttention 是 vLLM 提出的核心注意力算法,旨在解决大语言模型服务中的 KV Cache 内存管理低效问题。它借鉴了操作系统虚拟内存和分页的设计思想,将每个序列的 KV Cache 分割为固定大小的块(block),使得逻辑上连续的 KV 数据可以映射到物理上非连续的 GPU 内存块中 [S1][S29]。

在传统 LLM 推理框架中,KV Cache 通常以连续内存块的形式预先分配,大小为 max_model_len * per_token_kv_size。这种方式导致严重的内存浪费。PagedAttention 通过块级内存管理,将内存浪费降至接近零,同时支持灵活的请求间内存共享。

Block Table 数据结构

PagedAttention 的核心数据结构是 Block Table,它建立了从逻辑块到物理块的映射。在 vLLM V1 中,这一结构通过 BlockPool、KVCacheBlock、KVCacheManager 等类协同实现 [S4][S5][S7]。

KVCacheBlock(物理块元数据):

@dataclass(slots=True)
class KVCacheBlock:
    block_id: int                                    # 物理块 ID,范围 0 到 num_gpu_blocks - 1
    ref_cnt: int = 0                                 # 引用计数,表示有多少请求正在使用此块
    _block_hash: BlockHashWithGroupId | None = None  # 块内容哈希
    prev_free_block: "KVCacheBlock | None" = None    # 空闲队列双向链表前驱
    next_free_block: "KVCacheBlock | None" = None    # 空闲队列双向链表后继
    is_null: bool = False                            # 是否为 null_block(占位块,永不缓存)

ref_cnt 是引用计数机制的核心。当多个请求共享同一个物理块(如前缀缓存命中、并行采样共享 prompt)时,ref_cnt 递增;当请求释放块时,ref_cnt 递减。仅当 ref_cnt == 0 时,块才会被放回空闲队列。

块哈希计算(hash_block_tokens):

def hash_block_tokens(
    hash_function: Callable[[Any], bytes],
    parent_block_hash: BlockHash | None,
    curr_block_token_ids: Sequence[int],
    extra_keys: tuple[Any, ...] | None = None,
) -> BlockHash:
    if not parent_block_hash:
        parent_block_hash = NONE_HASH
    curr_block_token_ids_tuple = tuple(curr_block_token_ids)
    return BlockHash(
        hash_function((parent_block_hash, curr_block_token_ids_tuple, extra_keys))
    )

这里使用了父哈希链式结构:每个块的哈希不仅取决于自身包含的 token IDs,还取决于其父块(前一个块)的哈希。这确保了具有相同前缀的序列必然产生相同的块哈希序列,从而支持前缀缓存的精确匹配。

BlockPool(块池管理器):

BlockPool 是块级内存管理的核心 [S4],其主要职责包括:

  1. 分配新块: get_new_blocks(num_blocks) 从 free_block_queue 头部弹出指定数量的块,并将 ref_cnt 设为 1。
  2. 释放块: free_blocks(ordered_blocks) 递减 ref_cnt,将 ref_cnt == 0 的块追加回空闲队列。
  3. 触摸块(touch): touch(blocks) 将前缀缓存命中的块从空闲队列中移除(防止被驱逐),并递增 ref_cnt。
  4. 缓存完整块: cache_full_blocks() 将请求中已填满的块注册到 cached_block_hash_to_block 哈希表中。
  5. 驱逐缓存块: _maybe_evict_cached_block() 在分配新块时,若该块当前被缓存,则清除其哈希元数据并从缓存表中移除。

Copy-on-Write 机制

Copy-on-Write(CoW)是 PagedAttention 支持并行采样和束搜索的关键机制 [S1][S29]。当多个序列共享同一个物理块时,该块的 ref_cnt > 1。如果其中一个序列需要修改该块中的内容(如束搜索中某个分支对某个 token 的注意力输出不同),系统会创建一个该物理块的副本,而不是直接修改共享块。

在 vLLM V1 中,CoW 的实现细节体现在 allocate_slots 和 get_num_blocks_to_allocate 的交互中 [S6][S7]:

  1. 当请求需要写入新 token 到某个逻辑块时,KVCacheManager 检查该逻辑块对应的物理块是否被多个请求共享(ref_cnt > 1)。
  2. 如果是共享状态,系统从 BlockPool 分配一个新的物理块,将原块内容复制到新块,然后将原块的 ref_cnt 减 1。
  3. 请求的逻辑块指针更新为新分配的物理块,新块的 ref_cnt 设为 1。

这种块粒度的 CoW 比序列粒度的复制更高效:在束搜索中,只有发生分歧的块需要复制,共享前缀仍然保持共享。

非连续存储与内存分配策略

PagedAttention 的根本创新在于打破了 KV Cache 必须物理连续的假设:

  • 逻辑连续性: 每个请求的 KV Cache 在逻辑上仍是一个 token 序列,可通过逻辑块索引和块内偏移访问任意位置。
  • 物理非连续性: 相邻逻辑块可以映射到任意可用的物理块,由 BlockTable 维护映射关系。

在 vLLM V1 的 SingleTypeKVCacheManager 中,req_to_blocks 字典维护了每个请求的逻辑块到物理块的映射 [S6]:

self.req_to_blocks: defaultdict[str, list[KVCacheBlock]] = defaultdict(list)

当模型执行注意力计算时,attention backend(如 FlashAttention、FlashInfer)接收的是 block table(物理块 ID 列表)和 slot mapping(每个 token 对应的物理块内偏移),而非连续的内存指针。

内存计算公式

vLLM V1 中,PagedAttention 的 block_size 默认通常为 16 tokens。对于 70B 模型(FP16,80 层,8 KV heads,head_dim=128),每个 block 的内存约为:

block_memory = block_size * 2 * num_layers * num_kv_heads * head_dim * 2 bytes
             = 16 * 2 * 80 * 8 * 128 * 2
             = 16 * 2 * 80 * 8 * 256
             = 5,242,880 bytes ≈ 5.24 MB

在 80GB GPU 上,若 GPU 内存利用率为 0.9,预留约 72GB 给 KV Cache,则可容纳约 13,700 个 block,对应约 219,000 个 token 的并发容量。

与 OS 虚拟内存分页的类比和差异

维度OS 虚拟内存PagedAttention
页/块大小通常 4KB可配置,默认 16 tokens
逻辑-物理映射页表 (Page Table)Block Table (req_to_blocks)
物理内存组织非连续(分散的物理页框)非连续(分散的 GPU 内存块)
共享机制写时复制 (Copy-on-Write)块级 CoW (ref_cnt > 1 时复制)
缓存机制TLB (快表)BlockHashToBlockMap (前缀缓存)
置换策略LRU / ClockLRU (free_block_queue 顺序)
零页null pagenull_block (block_id=0)

关键差异:

  1. 块大小语义不同: OS 页大小是固定的内存字节数(如 4KB),而 PagedAttention 的 block_size 是 token 数量(如 16 tokens)。
  2. 无硬件 TLB: OS 有硬件 TLB 加速页表查询,PagedAttention 的 block table 查询由 GPU kernel 中的软件逻辑完成。
  3. 内容寻址缓存: OS 通常不使用内容寻址的页缓存,而 PagedAttention 的 prefix caching 本质上是内容寻址缓存。
  4. 分配粒度: OS 支持按需分页(demand paging),而 PagedAttention 目前是在调度阶段预分配块。

性能数据

根据 vLLM 原始论文 [S1]:

  • 吞吐量提升: 在多种模型(GPT、OPT、LLaMA)和工作负载下,vLLM 的吞吐量比 FasterTransformer 和 Orca 提升 2-4 倍。
  • 内存节省:
    • 并行采样(parallel sampling): 节省 6.1% - 9.8% 内存
    • 束搜索(beam search): 节省 37.6% - 55.2% 内存
  • 批处理规模: 由于内存碎片几乎消除,vLLM 可以支持更大的批处理规模,从而提升 GPU 利用率。

2.2 Continuous Batching

调度策略

Continuous Batching(持续批处理),又称 Iteration-level Scheduling(迭代级调度)或 Dynamic Batching(动态批处理),是一种 LLM 推理调度策略。与传统 Static Batching(静态批处理)不同,continuous batching 允许请求在每个模型前向传播迭代(iteration/step)的粒度上动态地加入或退出当前批处理,而不是等待整个批处理的所有请求都完成生成后才启动新批处理 [S2][S16][S17]。

这种调度方式的核心洞察是:LLM 的 decode 阶段每次只生成一个 token,不同请求的生成长度差异巨大。在静态批处理中,先完成的请求必须等待最后完成的请求,导致 GPU 在批次后期严重欠利用。

静态批处理的工作流程

  1. 收集一批请求,直到达到 batch_size 或超时。
  2. 对所有请求执行 prefill(计算 prompt 的 KV Cache)。
  3. 进入 decode 循环:每步为批处理中所有请求各生成一个 token。
  4. 当某个请求生成结束符(EOS)或达到最大长度时,该请求完成,但继续占用批处理槽位,用 padding 填充。
  5. 直到批处理中所有请求都完成后,才释放整个批处理,接收下一批请求。

问题分析:假设批处理中有 4 个请求,其生成长度分别为 100、200、300、400 tokens。在静态批处理中:

  • 第 100 步后,请求 1 已完成,但批处理继续运行 300 步,请求 1 的槽位被 padding 浪费。
  • 第 400 步后,整个批处理才结束。
  • 总 GPU 计算量中,padding 浪费的比例约为 (300×1 + 200×2 + 100×3) / (400×4) = 37.5%。

持续批处理的工作流程

  1. 每步开始时,调度器检查当前 running 请求的状态。
  2. 对于已完成的请求,立即从批处理中移除,释放其 KV Cache 槽位。
  3. 从 waiting 队列中选取新请求加入批处理(受限于 token budget 和 KV Cache 容量)。
  4. 新加入的请求可能处于 prefill 阶段(需要计算 prompt 的 KV),而现有请求处于 decode 阶段。
  5. 执行一步模型前向传播(可能混合 prefill 和 decode)。
  6. 重复步骤 1-5。

关键优势:

  • 零等待时间: 请求完成后立即被替换,不存在 padding 浪费。
  • 更高的 GPU 利用率: 批处理规模在每个 step 都保持较大,因为新请求不断补充。
  • 更优的延迟-吞吐量权衡: 不需要在 batch_size 和 latency 之间做痛苦的折中。

根据 Anyscale 的实测 [S17],continuous batching 相比 HuggingFace Transformers(静态批处理)可实现 23 倍吞吐量提升,p50 延迟降低。ORCA 论文 [S2] 报告在等效延迟目标下,相比 FasterTransformer 吞吐量提升 36.9 倍。

vLLM 调度器的三个队列

waiting 队列:等待被调度的请求队列。新到达的请求首先进入此队列。调度器按照 SchedulingPolicy(默认 FCFS,即先来先服务)排序。

skipped_waiting 队列:用于存放因异步依赖或约束而暂时无法调度的请求 [S3]。例如:

  • 请求正在等待远程 KV 传输完成(WAITING_FOR_REMOTE_KVS)
  • 请求正在等待结构化输出语法编译(WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR)
  • 请求正在等待流式输入(WAITING_FOR_STREAMING_REQ)

这些请求不会被调度执行,但会在每个 step 被检查是否可以"晋升"回 waiting 队列。

running 列表:当前正在执行的请求列表。注意这是一个 list[Request] 而非队列,因为调度器需要遍历它来决定每个 step 调度哪些请求。

每步调度的决策逻辑

Token Budget 与 Sequence Budget:

self.max_num_running_reqs = self.scheduler_config.max_num_seqs
self.max_num_scheduled_tokens = (
    self.scheduler_config.max_num_scheduled_tokens
    if self.scheduler_config.max_num_scheduled_tokens
    else self.scheduler_config.max_num_batched_tokens
)

调度受两个关键约束限制:

  1. max_num_running_reqs(最大并发请求数): 限制同时处于 running 状态的请求数量。
  2. max_num_scheduled_tokens(每步最大调度 token 数): 限制单个 step 中所有请求的总 token 数。

调度 RUNNING 请求:

调度器首先遍历 running 列表,为每个请求计算 num_new_tokens:

num_new_tokens = (
    request.num_tokens_with_spec
    + request.num_output_placeholders
    - request.num_computed_tokens
)
if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
    num_new_tokens = self.scheduler_config.long_prefill_token_threshold
num_new_tokens = min(num_new_tokens, token_budget)

对于每个请求,调度器尝试通过 kv_cache_manager.allocate_slots() 分配 KV Cache 块。若分配失败(内存不足),则触发抢占(preemption)。

抢占(Preemption)机制:

当 KV Cache 不足时,调度器选择优先级最低的 running 请求进行抢占 [S3]:

if self.policy == SchedulingPolicy.PRIORITY:
    preempted_req = max(
        self.running,
        key=lambda r: (r.priority, r.arrival_time),
    )
else:
    preempted_req = self.running.pop()

self._preempt_request(preempted_req, scheduled_timestamp)

被抢占的请求:

  • 状态从 RUNNING 变为 PREEMPTED
  • num_computed_tokens 重置为 0(丢弃已计算的 KV Cache)
  • 释放所有 KV Cache 块(kv_cache_manager.free(request))
  • 放回 waiting 队列头部(self.waiting.prepend_request(request))
  • num_preemptions 计数增加

值得注意的是,vLLM V1 目前采用的是重新计算(recomputation)策略而非换出(swap out)。即被抢占的请求不将 KV Cache 换出到 CPU 内存,而是直接丢弃,下次调度时重新计算。

调度 WAITING 请求:

在处理完所有 running 请求后,若还有剩余 token budget 且未达到 max_num_running_reqs,调度器开始处理 waiting 请求。对于 waiting 请求,调度器首先尝试前缀缓存命中:

new_computed_blocks, num_new_local_computed_tokens = (
    self.kv_cache_manager.get_computed_blocks(request)
)

若命中前缀缓存,则 num_computed_tokens 直接跳到缓存长度,无需重新计算前缀部分的 KV。

AsyncScheduler 异步调度机制

AsyncScheduler(vllm/v1/core/sched/async_scheduler.py)[S8] 继承自 Scheduler,实现了异步调度能力。其核心思想是:调度器在 CPU 上准备下一个 batch 的同时,GPU 正在执行当前 batch 的前向传播。这消除了 CPU 调度延迟对 GPU 利用率的影响。

AsyncScheduler 的关键修改在于 _update_after_schedule:

def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
    super()._update_after_schedule(scheduler_output)
    spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens
    for req_id in scheduler_output.num_scheduled_tokens:
        request = self.requests[req_id]
        if request.is_prefill_chunk:
            continue
        cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ()))
        request.num_output_placeholders += 1 + cur_num_spec_tokens
        request.spec_token_ids = self._spec_token_placeholders

这里引入了 num_output_placeholders 概念:对于每个被调度的 decode 请求,AsyncScheduler 预期它将在当前 step 生成 1 + num_spec_tokens 个新 token(1 个确定性的新 token + 若干个 speculative token)。这些 token 在调度时还未实际生成,因此用 placeholder 计数。

性能数据

指标静态批处理 (HF Transformers)Continuous Batching (vLLM)
相对吞吐量1x23x (Anyscale) / 24x (vLLM 论文)
等效延迟下吞吐量 vs FasterTransformer1x36.9x (ORCA)
GPU 利用率20-40%80-95%
批处理尾部 padding 浪费30-60%~0%

在实际生产环境中,continuous batching 的收益取决于工作负载的生成长度方差。方差越大(即请求间输出长度差异越大),continuous batching 相比 static batching 的优势越明显。

Continuous Batching 与 PagedAttention 的关系

Continuous Batching 依赖 PagedAttention 的灵活内存管理。在 static batching 中,每个请求的 KV Cache 槽位是固定的;而在 continuous batching 中,请求在每步结束后动态加入或退出,KV Cache 的分配和释放极其频繁。没有 PagedAttention 的细粒度 block 管理,这种动态调度将导致严重的内存碎片和分配失败。

PagedAttention 的 Copy-on-Write 机制使 continuous batching 中的并行采样和束搜索成为可能。当多个请求共享前缀时,它们可以共用相同的物理块,只在发生分歧时才复制,这大幅降低了 continuous batching 场景下的内存压力。

2.3 Chunked Prefill

性能优化

为什么需要 Chunked Prefill

在 LLM 推理中,Prefill 阶段和 Decode 阶段具有截然不同的计算特征:

  • Prefill: 计算密集型(compute-bound),一次性处理整个 prompt,可利用矩阵乘法的并行性。
  • Decode: 内存密集型(memory-bound),每次只处理一个(或少数几个)新 token,主要开销在于读取庞大的 KV Cache。

在没有 chunked prefill 的系统中,当一个长 prompt 请求(如 4096 tokens)进入系统时,它会在一个 step 中独占整个 GPU,执行 prefill 计算。在此期间:

  • 所有正在 decode 的请求必须等待,导致它们的 ITL(Inter-Token Latency) 急剧增加。
  • 若批处理中同时存在 prefill 和 decode 请求,整个 step 的执行时间由 prefill 决定(因为 prefill 的计算量远大于 decode),decode 请求被"阻塞"。

Chunked Prefill 通过将长 prefill 拆分为多个小块,使得 prefill 和 decode 可以交错执行,从而显著改善 ITL 和系统响应性 [S19][S25]。

Chunk 大小如何确定

Chunk 大小由 max_num_batched_tokens 参数控制 [S19][S25]。

在 vLLM V1 中,chunked prefill 默认开启。调度器在每个 step 中:

  1. 优先调度所有 pending decode 请求(因为 decode 请求的延迟敏感度高)。
  2. 将剩余 token budget 分配给 waiting 请求进行 prefill。
  3. 若某个 waiting 请求的 prefill token 数超过剩余 budget,则只调度 budget 内的部分,剩余部分留到下一个 step。

代码逻辑体现在 scheduler.py [S3]:

if (
    not self.scheduler_config.enable_chunked_prefill
    and num_new_tokens > token_budget
):
    break
num_new_tokens = min(num_new_tokens, token_budget)

max_num_batched_tokens 的调优:

  • 较小值(如 512-2048):
    • 优点:更好的 ITL,因为 prefill chunk 较小,不会长时间阻塞 decode 请求。
    • 缺点:更差的 TTFT(Time To First Token),因为长 prompt 需要更多 step 才能完成 prefill。
    • 适用场景:对流式输出、低延迟聊天应用敏感的场景。
  • 较大值(如 8192-32768):
    • 优点:更好的 TTFT 和吞吐量,因为更大的 chunk 可以更充分地利用 GPU 计算单元。
    • 缺点:更高的 ITL,因为大 chunk 的 prefill 会阻塞 decode 更长时间。
    • 适用场景:离线批处理、高吞吐量优先的场景。

vLLM V1 的默认值:

  • 在线服务(online serving): 8192
  • 离线处理(offline): 16384

AMD ROCm 文档建议 [S26]:

  • 交互式场景(低 TTFT): max_num_batched_tokens <= 8k-16k
  • 流式 decode(低 ITL): 16k-32k
  • 离线最大吞吐量: >= 32k

Prefill 和 Decode 的交错执行策略

在 vLLM V1 的调度策略中,decode 请求被优先调度 [S19]。具体流程:

  1. Decode 优先: 在 schedule() 的第一阶段,所有 running 请求(主要是 decode 阶段的请求)首先被考虑。它们的 num_new_tokens 通常为 1(或少量 speculative tokens)。
  2. Prefill 填充: 处理完所有 decode 请求后,剩余的 token_budget 用于调度 waiting 请求的 prefill。
  3. 混合执行: 一个 batch 中可能同时包含 prefill tokens(来自新请求的 chunk)和 decode tokens(来自现有请求)。vLLM V1 的 attention backend(如 FlashAttention)支持这种混合 batch 的注意力计算。

这种策略的好处是:

  • 计算与内存带宽重叠: Prefill 是计算密集的,decode 是内存带宽密集的。将它们放在同一个 batch 中,可以更均衡地利用 GPU 资源。
  • 延迟隐藏: 长 prefill 被拆分为多个 chunk,每个 chunk 只阻塞 decode 一小段时间。

vLLM V1 中默认开启的实现

vLLM V1 的 chunked prefill 实现是默认开启的,与 V0 需要显式设置 --enable-chunked-prefill 不同 [S19]。

在 Scheduler.__init__ 中 [S3]:

self.max_num_scheduled_tokens = (
    self.scheduler_config.max_num_scheduled_tokens
    if self.scheduler_config.max_num_scheduled_tokens
    else self.scheduler_config.max_num_batched_tokens
)

max_num_scheduled_tokens 实际上就是 max_num_batched_tokens。调度器在每个 step 中确保总调度 token 数不超过此限制。

对于 waiting 请求,long_prefill_token_threshold 参数提供了额外的限制 [S3]:

if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
    num_new_tokens = self.scheduler_config.long_prefill_token_threshold

这允许用户设置一个上限,防止单个请求在一个 step 中占用过多 token budget。

对 TTFT 和 ITL 的 Trade-off 分析

数学分析:

假设:

  • Prompt 长度: P
  • max_num_batched_tokens = B
  • Decode 阶段每个 step 生成 1 个 token

无 chunked prefill 时:

  • Prefill 在 1 个 step 内完成,耗时 T_prefill(P)。
  • TTFT = T_prefill(P)。
  • 对于正在 decode 的其他请求,该 step 中它们被完全阻塞,ITL 增加 T_prefill(P)。

有 chunked prefill 时:

  • Prefill 被拆分为 ceil(P / B) 个 chunk。
  • 每个 chunk 的耗时约为 T_prefill(B)(假设 prefill 时间与 token 数近似线性)。
  • TTFT = ceil(P / B) × T_step,其中 T_step 是每个 step 的总时间(包含交错 decode)。
  • 对于其他 decode 请求,每个 step 仅被阻塞 T_prefill(B),ITL 增加显著降低。

实际性能数据:

根据 vLLM 官方文档 [S19]:

  • 当 max_num_batched_tokens 从 512 提升到 8192 时:
    • TTFT 改善约 3-5 倍(因为更大的 chunk 减少了 prefill step 数)
    • ITL 轻微增加(约 10-20%)
    • 整体吞吐量提升约 20-40%

在 ShareGPT 等真实工作负载上,推荐的 max_num_batched_tokens 设置:

  • 小型模型(< 13B)+ 大 GPU(A100/H100): 16384 或更高
  • 大型模型(70B+): 8192-16384
  • 低延迟场景: 2048-4096

Chunked Prefill 与 Continuous Batching 的关系

Chunked Prefill 是 Continuous Batching 的增强。在纯 Continuous Batching 中,长 prefill 请求一旦开始就会阻塞所有 decode 请求直到完成。Chunked Prefill 通过将长 prefill 切分为小块,使得 prefill 和 decode 可以在同一个 step 中交错执行,从而在保持 high batch utilization 的同时,避免长 prefill 对 decode ITL 的剧烈冲击。

两者结合的工作流程:每个 step 中,调度器首先为所有 running decode 请求分配 token slot(每个 1 token),然后将剩余的 token budget 用于 waiting 请求的 prefill chunk。这样,decode 请求几乎不会被完全阻塞,而 prefill 请求也能逐步推进。

2.4 Prefix Caching

缓存优化

Prefix Caching(前缀缓存)是一种避免重复计算共享前缀 KV Cache 的优化技术。在 LLM 服务中,多个请求经常共享相同的前缀,例如:

  • 多轮对话中的系统 prompt 和历史记录
  • RAG(检索增强生成)中的相同检索上下文
  • 批量推理中使用相同的指令模板
  • 并行采样(parallel sampling)中相同的输入 prompt

Prefix Caching 的核心思想是:将已计算的前缀 KV Cache 存储在内存中,当新请求到达时,检查其前缀是否与已缓存的内容匹配。若匹配,则直接复用缓存的 KV Cache,只需计算不匹配的后缀部分。

vLLM 的 Block Hash 内容寻址机制

vLLM 的 Prefix Caching 实现基于 PagedAttention 的块级内容寻址 [S4][S5][S6]。

块哈希链:

每个块的哈希不仅取决于其包含的 token IDs,还取决于其父块的哈希。这形成了一个哈希链:

Block 0 hash = hash(NONE_HASH, tokens[0:BS], extra_keys_0)
Block 1 hash = hash(Block 0 hash, tokens[BS:2*BS], extra_keys_1)
Block 2 hash = hash(Block 1 hash, tokens[2*BS:3*BS], extra_keys_2)
...

其中 BS 是 block_size。这种链式结构确保了:

  • 两个请求若前 N 个 token 完全相同,则它们的前 N/BS 个块哈希必然完全相同。
  • 即使块内 token 相同,若父块不同,哈希也不同(防止不相关序列的碰撞)。

BlockHashToBlockMap:

BlockHashToBlockMap 是前缀缓存的核心索引结构 [S4]:

class BlockHashToBlockMap:
    def __init__(self):
        self._cache: dict[BlockHashWithGroupId, KVCacheBlock | dict[int, KVCacheBlock]] = {}

它支持三种操作:

  1. get_one_block(key): 返回具有给定哈希的任意一个物理块。
  2. insert(key, block): 将新块插入缓存。
  3. pop(key, block_id): 按块 ID 从缓存中移除特定块。

值得注意的是,vLLM 不进行块去重 [S4]:

NOTE #1: We currently don't de-duplicate the blocks in the cache,
meaning that if a block becomes full and is cached, we don't check
if there is already an identical block in the cache.

这是为了保持 block table 的追加-only 语义,简化调度器逻辑。

前缀缓存命中查找:

FullAttentionManager.find_longest_cache_hit [S6] 实现了最长前缀缓存命中查找:

@classmethod
def find_longest_cache_hit(cls, block_hashes, max_length, kv_cache_group_ids, block_pool, ...):
    computed_blocks = tuple([] for _ in range(len(kv_cache_group_ids)))
    block_size = kv_cache_spec.block_size
    max_num_blocks = max_length // block_size
    for block_hash in itertools.islice(block_hashes, max_num_blocks):
        if cached_block := block_pool.get_cached_block(block_hash, kv_cache_group_ids):
            for computed, cached in zip(computed_blocks, cached_block):
                computed.append(cached)
        else:
            break
    if use_eagle and computed_blocks[0]:
        for computed in computed_blocks:
            computed.pop()  # EAGLE 需要丢弃最后一个匹配块
    return computed_blocks

查找逻辑是顺序扫描:从第一个块开始,依次检查每个块哈希是否在缓存中。一旦遇到未缓存的块,立即停止(因为哈希链的性质保证了后续块必然也未缓存)。

块缓存与驱逐:

当请求完成一个块的计算后,cache_blocks 方法将其注册到缓存 [S6][S7]:

def cache_blocks(self, request: Request, num_tokens: int) -> None:
    num_cached_blocks = self.num_cached_block.get(request.request_id, 0)
    num_full_blocks = num_tokens // self.block_size
    if num_cached_blocks >= num_full_blocks:
        return
    self.block_pool.cache_full_blocks(
        request=request,
        blocks=self.req_to_blocks[request.request_id],
        num_cached_blocks=num_cached_blocks,
        num_full_blocks=num_full_blocks,
        block_size=self.block_size,
        kv_cache_group_id=self.kv_cache_group_id,
    )
    self.num_cached_block[request.request_id] = num_full_blocks

只有完整填满的块(即包含 block_size 个 token)才会被缓存。这是为了避免缓存不完整的块,因为不完整块的内容可能在后续步骤中改变。

SGLang 的 RadixAttention(Radix Tree)对比

特性vLLM (PagedAttention + APC)SGLang (RadixAttention)
数据结构哈希表 (BlockHashToBlockMap)Radix Tree(压缩 Trie)
匹配粒度块级 (block_size tokens)令牌级 (token-level)
匹配方式精确块哈希匹配最长公共前缀匹配
分支支持有限(块表追加-only)原生支持(树的分支结构)
配置需求需显式开启 --enable-prefix-caching自动,零配置
驱逐策略LRU(free_block_queue 顺序)LRU(叶节点)
生成结果缓存不支持支持(多轮历史复用)

Radix Tree 工作原理:

RadixAttention 将 KV Cache 块组织为一棵 Radix Tree(基数树/压缩前缀树)[S24]:

  • 边(Edge): 存储 token 序列(而非单个 token)。
  • 节点(Node): 对应一个 KV Cache 块。
  • 根到叶路径: 表示一个完整的 token 序列。

当新请求到达时:

  1. 从根节点开始,沿树向下遍历。
  2. 将请求的 token 序列与边上的 token 序列逐一匹配。
  3. 找到最长公共前缀对应的节点。
  4. 若在某个边上部分匹配,则将该边分裂为两条:公共前缀部分和新后缀部分。
  5. 从匹配点继续计算剩余的 token。

关键差异分析:

vLLM 的优势:

  • 块级哈希实现简单,查找速度快(O(num_blocks))。
  • 与 PagedAttention 的块管理天然契合。
  • 工业验证充分,vLLM 已被广泛部署。

vLLM 的劣势:

  • 匹配粒度受限于 block_size(默认 16 tokens)。若前缀长度不是 block_size 的整数倍,最后一块无法命中。
  • 不支持动态分支。一旦某请求修改了共享块中的内容(如 beam search 中的分支),需要 CoW 复制整个块。

SGLang 的优势:

  • 令牌级匹配精度更高,不受 block_size 限制。
  • 天然支持分支和树形结构(如多轮对话、agent 工具调用)。
  • 自动发现缓存机会,无需配置。

SGLang 的劣势:

  • Radix Tree 的遍历和维护开销高于哈希表。
  • 在高度随机的请求分布下,Radix Tree 的缓存命中率可能不如块级哈希稳定。

实际性能对比 [S22]:

  • 在 80% 共享前缀的工作负载下,SGLang 的 TTFT 比 vLLM 快 2-6 倍。
  • 在无共享前缀的随机工作负载下,两者性能差异缩小到 10-20%。

自动前缀缓存的命中条件和失效策略

命中条件:

在 vLLM V1 中,前缀缓存命中需要满足以下条件 [S6][S7]:

  1. 块满: 只有包含完整 block_size 个 token 的块才会被缓存。
  2. 哈希一致: 块的哈希必须与请求当前位置的预期哈希完全匹配(包括父块哈希链和 extra_keys)。
  3. KV Cache Group 一致: 对于混合模型,每个 KV cache group 必须独立命中。
  4. EAGLE 特殊处理: 若启用 EAGLE,最后一个匹配块会被丢弃,因为 EAGLE 需要该块的隐藏状态来生成 draft token。
  5. 滑动窗口约束: 对于 sliding window attention,需要至少 sliding_window_contiguous_blocks 个连续命中块。

失效策略:

vLLM V1 的前缀缓存失效策略包括 [S4][S7]:

  1. LRU 驱逐: 当需要分配新块且空闲队列耗尽时,从队列头部(最久未使用)取出的块若带有缓存哈希,则自动驱逐。
  2. 显式重置: reset_prefix_cache() 方法可清空所有缓存。
  3. 滑动窗口回收: 对于 sliding window attention,超出窗口范围的块被自动释放。
  4. 请求完成释放: 请求完成后,free_blocks 递减 ref_cnt。当 ref_cnt 降为 0 时,块回到空闲队列,但不自动清除其缓存哈希。只有当该块被重新分配时,其缓存哈希才会被驱逐。

多轮对话和 RAG 场景中的收益

多轮对话:

在多轮对话中,后续轮次通常共享前几轮的对话历史作为前缀。Prefix Caching 的收益取决于:

  • 命中长度: 若每轮新增内容较少(如 50 tokens),而历史较长(如 2000 tokens),则缓存命中率可达 2000 / 2050 ≈ 97.5%。
  • 块对齐: 由于 vLLM 按块缓存,若历史长度不是 block_size 的倍数,最后一块(最多 block_size-1 个 token)需要重新计算。
  • 动态变化: 若对话历史在不同请求间有细微差异(如时间戳、用户名),前缀缓存可能完全失效。此时 SGLang 的 RadixAttention 更具优势,因为它可以在树的不同分支上共享公共前缀。

RAG 场景:

在 RAG 中,多个请求可能使用相同的检索文档作为上下文:

  • 文档级共享: 若多个查询使用相同的检索结果,文档部分的 KV Cache 可完全共享。
  • 查询级差异: 每个查询的问题部分不同,但从文档末尾开始分叉。vLLM 的块级缓存可以有效共享文档块,而查询块独立计算。
  • 收益估算: 假设文档长度为 1024 tokens,查询长度为 64 tokens,block_size=16。则文档部分有 64 个块可共享,查询部分有 4 个块独立计算。Prefill 计算量减少约 64/68 ≈ 94%。

性能数据:

场景前缀共享比例TTFT 降低吞吐量提升
多轮聊天(系统 prompt 共享)60-80%40-60%1.5-2x
RAG(相同文档)80-95%60-80%2-4x
批量模板推理90-99%80-90%3-6x
无共享前缀0%0%无变化

Prefix Caching 与 PagedAttention 的关系

Prefix Caching 建立在 PagedAttention 的 block 管理之上。没有 PagedAttention 的细粒度 block 分配和引用计数机制,Prefix Caching 无法实现:

  • 多个请求共享同一个物理块的 KV Cache(通过 ref_cnt 管理)。
  • 块级别的内容寻址(通过 _block_hash 字段)。
  • 缓存块的按需分配和驱逐(通过 BlockPool 的空闲队列管理)。

PagedAttention 的 Copy-on-Write 机制是 Prefix Caching 安全共享的前提。当多个请求共享缓存的前缀块时,若其中一个请求需要修改该块(如生成了新的内容),CoW 机制会自动复制该块,确保其他请求仍能安全地读取原始缓存内容。

2.5 Speculative Decoding

加速技术

Speculative Decoding(投机解码)是一种加速 LLM 自回归生成的技术。其核心思想是:使用一个快速的 Draft Model(草稿模型)预测未来若干个 token,然后用慢速但精确的 Target Model(目标模型)并行验证这些预测。验证过程中,目标模型一次性处理 draft tokens 序列,通过概率比较决定接受或拒绝每个 draft token。被接受的 draft tokens 直接输出,无需重新计算;被拒绝位置之后的 tokens 被丢弃,从拒绝点重新开始 [S21][S27]。

投机解码的关键洞察是:在 GPU 上执行一次大规模前向传播(验证多个 token)的速度远快于执行多次小规模前向传播(逐个生成 token)。即使 draft model 的准确率不高,只要验证步骤足够快,整体吞吐量仍可显著提升。

Draft 方式对比

方法原理额外内存开销加速比典型接受长度适用场景
N-gram基于输入序列本身的 n-gram 匹配生成 draft token无1.0-1.3x1.1-1.5重复性高的文本(代码、模板)
Draft Model(独立小模型)使用完整独立小模型(如 7B draft 70B target)~10%1.5-2.0x2-3通用场景
EAGLE在 target model 隐藏状态之上训练的轻量级外推头(1-2 层 Transformer)~3-5%1.54-2.36x2.5-3.5高质量 draft 需求
MLP Draft(Medusa)使用简单 MLP 层从 target 隐藏状态预测未来 token很小1.2-1.5x1.5-2快速部署
MTP(Multi-Token Prediction)将投机解码能力内建到目标模型中(如 DeepSeek V3)内建1.5-1.8x2-3模型原生支持

N-gram 方法

NgramProposer 不需要任何额外模型,完全基于输入序列本身的 n-gram 匹配来生成 draft token。实现上使用 Numba JIT 编译加速的 LPS(Longest Prefix Suffix)算法 [S11]:

@jit(nopython=True)
def _find_longest_matched_ngram_and_propose_tokens(origin_tokens, min_ngram, max_ngram, ...):
    tokens = origin_tokens[::-1]
    lps = np.zeros(max_ngram, dtype=np.int32)
    longest_ngram = 0
    position = 0
    prev_lps = 0
    i = 1
    while i < total_token:
        if tokens[prev_lps] == tokens[i]:
            prev_lps += 1
            if prev_lps >= longest_ngram:
                longest_ngram = prev_lps
                position = i
            if i < max_ngram:
                lps[i] = prev_lps
            if prev_lps == max_ngram:
                prev_lps = lps[max_ngram - 1]
            i += 1
        elif prev_lps != 0:
            prev_lps = lps[prev_lps - 1]
        else:
            i += 1

优势:零额外内存开销,对重复性高的文本效果极佳。

劣势:对创造性文本效果差。

EAGLE 方法

EagleProposer 使用 EAGLE 方法。EAGLE 不是完整的独立模型,而是在 target model 的隐藏状态之上训练的一个轻量级外推头(通常只有 1-2 层 Transformer)。

class EagleProposer(SpecDecodeBaseProposer):
    def __init__(self, vllm_config, device, runner=None):
        super().__init__(vllm_config, device,
                         pass_hidden_states_to_model=True, runner=runner)

pass_hidden_states_to_model=True 表示 EAGLE 需要接收 target model 的隐藏状态作为输入,而非原始 token IDs。

EAGLE 的优势:

  • 模型体积极小(通常只有 target 模型的 1-5%)。
  • 直接利用 target 的特征表示,draft 质量高。
  • 可与 target 模型共享部分 KV Cache。

SGLang 的测试显示 [S27],EAGLE-2 在 LLaMA-3.1-8B 上可达 1.54x 加速,EAGLE-3 可达 2.36x 加速。

原生 Multi-Token Prediction(MTP)

Multi-Token Prediction(MTP)是一种将投机解码能力内建到目标模型中的技术 [S21][S24][S28]。与需要独立 draft model 的方法不同,MTP 在目标模型的训练阶段就加入了多个预测头,每个头负责预测未来不同位置的 token。

DeepSeek V3 的 MTP:

DeepSeek V3 采用了 MTP 架构 [S21]。其结构特点:

  • 在主模型的每个 Transformer 层(或顶层)附加多个预测头。
  • 第 i 个预测头基于当前隐藏状态预测第 i+1 个未来 token。
  • 所有预测头共享主模型的 KV Cache。

在推理时:

  1. 主模型执行一次前向传播。
  2. 各 MTP 头并行输出多个候选 token。
  3. 主模型的最终输出层验证这些候选 token。

Gemma 4 的 Assistant Models:

Gemma 4 使用了专门的 Assistant Model(助手模型)实现 MTP [S20][S28]。vLLM 中的 Gemma4Proposer [S13] 专门处理这种情况:

class Gemma4Proposer(SpecDecodeBaseProposer):
    def __init__(self, vllm_config, device, runner=None):
        super().__init__(vllm_config, device,
                         pass_hidden_states_to_model=True, runner=runner)
        self.constant_draft_positions = True

constant_draft_positions = True 是 Gemma 4 MTP 的关键特征:所有 draft 步骤从同一位置预测(即 target model 的最后一个位置),而不是像自回归那样逐步前进。这是因为 Gemma 4 的 assistant model 使用 cross-attention 或 Q-only attention 机制,直接基于 target 的完整上下文预测未来 token。

Gemma4Proposer 还实现了与 target model 的 KV Cache 共享 [S13]:

def _setup_gemma4_kv_sharing(self, target_attn_layer_names):
    for draft_idx, layer in enumerate(self.model.model.layers):
        attn = getattr(layer.self_attn, "attn", None)
        draft_layer_type = draft_layer_types[draft_idx]
        candidates = type_to_target_indices.get(draft_layer_type, [])
        target_idx = candidates[-1]
        target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
        attn.kv_sharing_target_layer_name = target_layer_name

每个 draft decoder layer 映射到相同 attention 类型的最后一个非共享 target layer,从而实现 KV Cache 的跨模型复用。

验证流程的详细步骤与接受率计算

标准验证流程:

投机解码的验证流程如下 [S12][S27]:

  1. Draft 阶段: Draft model 生成 k 个候选 token:d₁, d₂, ..., dₖ。
  2. Target 前向传播: Target model 接收原始序列 [x₁, ..., xₜ] 和 draft tokens [d₁, ..., dₖ],一次性计算 k+1 个位置的概率分布:
    • p₀ = P(target | x₁, ..., xₜ)(原始下一个 token 的分布)
    • p₁ = P(target | x₁, ..., xₜ, d₁)
    • ...
    • pₖ = P(target | x₁, ..., xₜ, d₁, ..., dₖ)
  3. 逐位置接受/拒绝:
    • 对于位置 i,计算 draft model 在该位置的概率 qᵢ = P(draft | context, d₁, ..., dᵢ₋₁)。
    • 接受概率为 min(1, pᵢ(dᵢ) / qᵢ(dᵢ))。
    • 以该概率接受 dᵢ。若接受,继续验证 dᵢ₊₁。
    • 若拒绝,从修正分布 norm(max(0, pᵢ − qᵢ)) 中采样一个新 token x',输出 x',并丢弃 dᵢ₊₁, ..., dₖ。
  4. Bonus Token: 若所有 k 个 draft token 都被接受,target model 在最后一个位置已经计算了 pₖ,可以直接从中采样一个额外的 token,相当于一次验证获得了 k+1 个 token。

接受率计算:

接受率(Acceptance Rate)是衡量投机解码效果的核心指标 [S14]:

@dataclass
class SpecDecodingStats:
    num_spec_tokens: int
    num_drafts: int = 0
    num_draft_tokens: int = 0
    num_accepted_tokens: int = 0
    num_accepted_tokens_per_pos: list[int] = field(default_factory=list)

    def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int):
        self.num_drafts += 1
        self.num_draft_tokens += num_draft_tokens
        self.num_accepted_tokens += num_accepted_tokens

关键指标:

  • Draft Acceptance Rate = num_accepted_tokens / num_draft_tokens × 100%
  • Mean Acceptance Length = 1 + (num_accepted_tokens / num_drafts)(包含 bonus token)
  • Per-position Acceptance Rate: 第 i 个 draft 位置被接受的概率,通常随位置增加而递减。

在 vLLM V1 的 Scheduler.update_from_output [S3] 中,验证结果的统计如下:

if scheduled_spec_token_ids and generated_token_ids:
    num_draft_tokens = len(scheduled_spec_token_ids)
    num_accepted = len(generated_token_ids) - 1
    num_rejected = num_draft_tokens - num_accepted
    if request.num_computed_tokens > 0:
        request.num_computed_tokens -= num_rejected

与 Continuous Batching 的兼容性(Ragged Tensor)

投机解码与 continuous batching 的结合面临一个关键挑战:Ragged Tensor(不规则张量)[S12]。

在标准 continuous batching 中,每个请求每步只生成 1 个 token,batch 中所有请求的序列长度增量相同,便于组织张量。但在投机解码中:

  • 不同请求可能生成不同数量的 draft tokens(如 0 到 num_speculative_tokens)。
  • 验证后,不同请求实际接受的 token 数也不同。
  • 这导致 batch 中各请求的序列长度不一致,形成 "ragged" 结构。

vLLM V1 通过 Padded Drafter Batch 解决这个问题 [S12][S15]:

self.disable_padded_drafter_batch: bool = False
"""Disable input padding for speculative decoding. If set to True,
speculative input batches can contain sequences of different lengths,
which may only be supported by certain attention backends."""

默认情况下,vLLM 使用 padding 方式:将所有请求扩展到 batch 中最长序列的长度,用特殊 token(如 padding token 或 rejected token mask)填充。这虽然增加了少量无效计算,但保持了张量的规则性,便于与各种 attention backend 兼容。

vLLM 的 speculative_config 配置

SpeculativeConfig(vllm/config/speculative.py)[S15] 是投机解码的配置中心:

@config
class SpeculativeConfig:
    num_speculative_tokens: int = Field(default=None, gt=0)
    model: str | None = None
    method: SpeculativeMethod | None = None
    draft_tensor_parallel_size: int | None = Field(default=None, ge=1)
    prompt_lookup_max: int | None = Field(default=None, ge=1)
    prompt_lookup_min: int | None = Field(default=None, ge=1)
    parallel_drafting: bool = False
    disable_padded_drafter_batch: bool = False
    use_local_argmax_reduction: bool = False
    rejection_sample_method: RejectionSampleMethod = "standard"

关键配置项:

  • num_speculative_tokens: 每次 draft 生成的 token 数。值越大,潜在加速越高,但 draft 错误率也可能增加。
  • method: 投机方法,可选值包括 "ngram", "medusa", "mlp_speculator", "draft_model", "eagle", "eagle3", "mtp", "dflash" 等。
  • model: draft 模型路径(用于 draft_model、EAGLE、MTP 等方法)。
  • parallel_drafting: 是否启用并行 draft。

Speculative Decoding 与其他概念的关系

与 Continuous Batching 的关系:投机解码增加了 batch 中序列长度的方差,对 continuous batching 的调度提出了更高要求。AsyncScheduler 的 num_output_placeholders 机制专门用于管理 speculative decode tokens 的调度占位。

与 PagedAttention 的关系:每接受一个 draft token,对应请求的 block table 和 seq_len 需要更新。eagle_step_update_slot_mapping_and_metadata Triton kernel 负责此操作 [S12]。EAGLE 和 Gemma 4 MTP 等方法还需要 draft model 与 target model 共享 KV Cache。

与 Chunked Prefill 的关系:在长 prefill 场景中,投机解码的收益可能降低,因为 prefill 阶段本身已经是计算密集的,decode 阶段的比例相对较小。

← 上一模块
LLM 推理基础概念
下一模块 →
vLLM API 模型类型