大模型推理加速：从 KV Cache 到投机解码的工程实践

baronbool

10400人浏览 · 2026-06-08 08:19:06

baronbool · 2026-06-08 08:19:06 发布

大模型推理加速：从 KV Cache 到投机解码的工程实践

cover

一、推理延迟的"最后一公里"：模型能力够，但速度不够

大模型应用落地的最大瓶颈往往不是模型能力，而是推理延迟。一个 70B 参数的模型，单次生成可能需要 5-10 秒，在实时对话和批量处理场景中都难以接受。用户对响应时间的容忍度通常在 2 秒以内，超过 5 秒就会产生明显的体验下降。

推理加速的核心矛盾在于：自回归生成的串行性——每个 token 的生成都依赖前面所有 token 的计算结果，无法简单并行化。解决这一矛盾，需要从计算复用（KV Cache）、计算投机（投机解码）、模型压缩（量化）三个维度同时发力。

二、推理加速技术栈全景

graph TB
    subgraph 计算复用层
        A[KV Cache<br/>避免重复计算Key/Value]
        A1[PagedAttention<br/>显存分页管理]
    end

    subgraph 计算投机层
        B[投机解码<br/>小模型猜+大模型验]
        B1[Medusa Heads<br/>多头并行预测]
    end

    subgraph 模型压缩层
        C[量化<br/>INT8/INT4/FP8]
        C1[蒸馏<br/>知识迁移到小模型]
    end

    subgraph 系统优化层
        D[Continuous Batching<br/>动态批处理]
        D1[Prefix Caching<br/>共享前缀缓存]
    end

    A --> A1
    B --> B1
    C --> C1
    D --> D1

    A1 --> E[端到端推理加速]
    B1 --> E
    C1 --> E
    D1 --> E

KV Cache 是最基础也最有效的优化。自回归生成中，每一步都需要对前面所有 token 计算 Attention，而 Key 和 Value 矩阵在之前步骤已经计算过。KV Cache 将这些矩阵缓存下来，避免重复计算，将每步的计算量从 O(n²) 降低到 O(n)。

三、核心加速方案实现

3.1 KV Cache 与 PagedAttention

import torch
from dataclasses import dataclass

@dataclass
class KVCache:
    """KV Cache 管理：支持 PagedAttention 的分页存储"""
    key_cache: torch.Tensor    # [num_layers, max_pages, page_size, num_heads, head_dim]
    value_cache: torch.Tensor  # 同上
    page_table: dict[int, list[int]]  # seq_id -> 页号列表
    free_pages: list[int]      # 空闲页列表

class PagedKVCacheManager:
    """分页 KV Cache：解决显存碎片问题"""

    def __init__(
        self,
        num_layers: int,
        num_heads: int,
        head_dim: int,
        page_size: int = 16,
        max_pages: int = 4096,
        dtype: torch.dtype = torch.float16
    ):
        self.page_size = page_size
        # 预分配显存，避免动态分配的碎片问题
        cache_shape = (num_layers, max_pages, page_size, num_heads, head_dim)
        self.key_cache = torch.zeros(cache_shape, dtype=dtype, device="cuda")
        self.value_cache = torch.zeros(cache_shape, dtype=dtype, device="cuda")
        self.page_table = {}
        self.free_pages = list(range(max_pages))

    def allocate(self, seq_id: int, num_tokens: int) -> list[int]:
        """为序列分配 KV Cache 页"""
        num_pages = (num_tokens + self.page_size - 1) // self.page_size
        if len(self.free_pages) < num_pages:
            # 显存不足，触发驱逐策略
            self._evict(num_pages - len(self.free_pages))

        pages = self.free_pages[:num_pages]
        self.free_pages = self.free_pages[num_pages:]
        self.page_table[seq_id] = pages
        return pages

    def _evict(self, needed: int) -> None:
        """LRU 驱逐：释放最久未访问的序列的页"""
        # 按最后访问时间排序，驱逐最早的
        sorted_seqs = sorted(
            self.page_table.items(),
            key=lambda x: x[0]  # 简化：实际应按访问时间
        )
        evicted = 0
        for seq_id, pages in sorted_seqs:
            self.free_pages.extend(pages)
            del self.page_table[seq_id]
            evicted += len(pages)
            if evicted >= needed:
                break

3.2 投机解码

class SpeculativeDecoder:
    """投机解码：小模型猜测 + 大模型验证"""

    def __init__(self, draft_model, target_model, gamma: int = 5):
        self.draft_model = draft_model   # 小模型（快速）
        self.target_model = target_model  # 大模型（准确）
        self.gamma = gamma               # 猜测token数

    def generate(self, prompt_ids: list[int], max_tokens: int) -> list[int]:
        generated = list(prompt_ids)

        while len(generated) - len(prompt_ids) < max_tokens:
            # 步骤1：小模型快速生成 gamma 个 token
            draft_tokens = []
            draft_probs = []
            current = list(generated)

            for _ in range(self.gamma):
                next_token, prob = self.draft_model.predict_next(current)
                draft_tokens.append(next_token)
                draft_probs.append(prob)
                current.append(next_token)

            # 步骤2：大模型一次前向传播验证所有猜测token
            target_probs = self.target_model.predict_batch(
                generated, draft_tokens
            )

            # 步骤3：逐个验证，接受或拒绝
            accepted = 0
            for i in range(self.gamma):
                # 接受条件：随机采样接受概率
                # p_accept = min(1, target_prob / draft_prob)
                t_prob = target_probs[i][draft_tokens[i]]
                d_prob = draft_probs[i][draft_tokens[i]]
                accept_ratio = min(1.0, t_prob / (d_prob + 1e-10))

                if torch.rand(1).item() < accept_ratio:
                    generated.append(draft_tokens[i])
                    accepted += 1
                else:
                    # 拒绝后，从大模型的分布中采样一个token
                    corrected = torch.multinomial(
                        target_probs[i], num_samples=1
                    ).item()
                    generated.append(corrected)
                    break
            else:
                # 所有猜测都被接受，额外生成一个token
                extra = torch.multinomial(
                    target_probs[self.gamma], num_samples=1
                ).item()
                generated.append(extra)

        return generated[:len(prompt_ids) + max_tokens]