AI 推理引擎的批处理与内存管理：从 Continuous Batching 到 PagedAttention 的工程实践

一、引言：推理引擎的"吞吐-延迟"不可能三角

当你部署一个 70B 参数的大语言模型时，会面临一个根本矛盾：

• 吞吐（Throughput）：每秒处理尽可能多的 token，摊薄 GPU 成本
• 延迟（Latency）：单请求的首 token 时间（TTFT）和流式输出间隔要尽可能低
• 显存（Memory）：KV Cache 占用随序列长度线性增长，单卡 HBM 有限

传统做法是静态批处理（Static Batching）：凑齐一批请求一起执行，等最慢的请求完成才能释放资源。这在推理引擎中导致了严重的"队头阻塞"——一个长序列请求会拖慢同批所有短序列请求。

本文深入剖析现代推理引擎解决这一问题的两大核心技术：Continuous Batching（迭代级批处理） 和 PagedAttention（分页 KV Cache），从操作系统内存管理的视角理解其工程原理。

二、静态批处理：为什么 GPU 利用率上不去

2.1 Padding 与资源浪费

静态批处理需要将所有序列 pad 到相同长度。假设一个 batch 中有： - 请求 A：输入 10 token，输出 200 token - 请求 B：输入 500 token，输出 5 token - 请求 C：输入 50 token，输出 50 token

静态批处理必须等 A 生成完 200 token 才能释放整个 batch。B 仅需 5 token 却被迫等待 195 个无用迭代。更糟的是，每次迭代的 attention 计算都需要 pad 到最大长度，造成大量无效 FLOPs。

2.2 KV Cache 的显存碎片

每个请求的 KV Cache 是 [seq_len, num_heads, head_dim] 的三维张量。静态分配时必须预留最大序列长度（如 4096 或 8192 token）。实际推理时： - 有的序列只用了 200 token，剩余 3800 token 的空间被闲置 - 不同请求的序列长度不同，显存池中产生大量无法被复用的小碎片

这就是经典的内存外部碎片问题——操作系统课程中分页机制要解决的正是一类问题。

三、Continuous Batching：事件驱动与迭代级调度

3.1 核心思想

Continuous Batching（又称 Iteration-Level Batching）将调度粒度从"请求级"细化到"迭代级"：

1. 每次模型 forward 之前，Scheduler 先检查当前 iteration
2. 新到达的请求立即加入 batch
3. 已完成（生成 EOS token）的请求立即移出 batch
4. 超出显存容量的请求被 preempted（抢占），swap 到 CPU 内存

这与操作系统中的抢占式多任务调度如出一辙——每个 iteration 是一个时间片，Scheduler 每次重新决定谁上谁下。

3.2 vLLM 的实现：Preemption 机制

vLLM（斯坦福大学开源的高性能推理引擎）中的 Preemption 策略：

# 简化示意：当 GPU KV Cache 不足时的抢占逻辑
class Scheduler:
    def schedule(self) -> SchedulerOutputs:
        # 处理抢占：从 running 队列尾部回收资源
        while self.running and self.kv_cache_manager.available_blocks < needed_blocks:
            preempted_req = self.running.pop()  # 抢占最后一个请求
            self._preempt_request(preempted_req)
            self.waiting.appendleft(preempted_req)  # 放回等待队列头部

        # 从 waiting 队列调度新请求
        scheduled = []
        while self.waiting and self.kv_cache_manager.can_allocate(needed_blocks):
            req = self.waiting.popleft()
            self._allocate_kv_blocks(req)
            scheduled.append(req)

        return SchedulerOutputs(scheduled, ...)

    def _preempt_request(self, request):
        """恢复被抢占请求的 KV Cache"""
        if self.num_preemption >= MAX_PREEMPTIONS:
            # 超过抢占阈值，执行 recomputation 而非 swap
            self._deallocate_kv_blocks(request)
        else:
            # 第一次抢占：将 KV Cache swap 到 CPU 内存
            self._swap_to_cpu(request)
            self.num_preemption += 1

vLLM 采用两种抢占策略： - Swapping：将 KV Cache 从 GPU HBM 搬到 CPU DRAM - Recomputation：丢弃 KV Cache，后续重新计算

Swapping 适合 KV Cache 较大的情况（如长序列），Recomputation 适合短序列或 prefill 阶段。

3.3 调度策略：FCFS 与 WaterLevel

除了基本的 FCFS（先来先服务），Modern TensorRT-LLM 和 SGLang 还支持：

• WaterLevel Strategy：设置 KV Cache 的低水位（Low Watermark）和高水位（High Watermark）。水位低于下限时立即抢占，低于上限时拒绝新请求。
• Chunked Prefill：将长序列的 prefill 阶段切分为多个 chunk，与 decode 交错执行，避免单次 prefill 阻塞所有 decode 请求。

四、PagedAttention：操作系统的分页内存管理照进 GPU

4.1 核心洞察：GPU 上的虚拟内存分页

PagedAttention 的灵感直接来自操作系统的虚拟内存系统：

• 逻辑页（Logical Page） → 序列的连续 token 块（如每块 16 token）
• 物理页（Physical Page） → GPU HBM 中的 KV Cache block
• 页表（Page Table） → 从逻辑块索引到物理块地址的映射

在标准推理引擎中，KV Cache 是连续分配的：

传统方式: |██████████████████████████████|  连续 4096 token 的物理内存
                    ↑
              实际只用 200 token，其余闲置

PagedAttention 将其打散为固定大小的 block：

PagedAttention:  逻辑视图 → [Block 0] [Block 1] [Block 2] ... [Block N]
                          ↓           ↓           ↓              ↓
                 物理内存 → [GPU #17] [GPU #42] [GPU #08] ... [GPU #91]

                 Page Table: {0→17, 1→42, 2→8, ..., N→91}

不同请求的 block 可以共享（用于 Beam Search 和 Prefix Caching），实现写时复制（Copy-on-Write）——与操作系统 fork 共享物理页完全一致。

4.2 Block Table 的数据结构

在 vLLM 中，每个序列维护一个 Block Table：

# vLLM 核心数据结构（简化）
@dataclass
class Sequence:
    seq_id: int
    block_size: int = 16          # 每个 block 存储的 KV token 数
    block_table: List[int] = []   # 逻辑块 → 物理块 的映射表
    prompt_token_ids: List[int] = []
    output_token_ids: List[int] = []

    def append_token(self, token_id):
        """追加一个 token，必要时分配新 block"""
        self.output_token_ids.append(token_id)
        total_tokens = len(self.prompt_token_ids) + len(self.output_token_ids)

        # 计算需要的总 block 数
        needed_blocks = ceil(total_tokens / self.block_size)

        if len(self.block_table) < needed_blocks:
            # 分配新物理 block
            new_block = self.allocator.allocate()
            self.block_table.append(new_block)

    def get_physical_blocks(self) -> List[int]:
        return self.block_table.copy()

    def free_blocks(self):
        """序列结束时，释放所有物理 block 到空闲池"""
        for block in self.block_table:
            self.allocator.free(block)
        self.block_table.clear()

Block Allocator 维护一个物理 block 的 LRU 空闲列表和引用计数：

class BlockAllocator:
    def __init__(self, num_gpu_blocks: int, block_size: int = 16):
        self.block_size = block_size
        self.free_blocks: Deque[int] = deque(range(num_gpu_blocks))
        self.ref_counts: Dict[int, int] = defaultdict(int)  # 引用计数（用于 COW）

    def allocate(self) -> int:
        if not self.free_blocks:
            raise MemoryError("GPU KV Cache exhausted")
        block_id = self.free_blocks.popleft()
        self.ref_counts[block_id] = 1
        return block_id

    def free(self, block_id: int):
        self.ref_counts[block_id] -= 1
        if self.ref_counts[block_id] == 0:
            self.free_blocks.append(block_id)

    def cow_duplicate(self, block_id: int) -> int:
        """写时复制：仅当多序列共享该 block 时，实际执行拷贝"""
        if self.ref_counts[block_id] == 1:
            return block_id  # 唯一持有者，无需拷贝
        # 需要分裂：拷贝原 block，旧 block 引用减 1，新 block 开始
        new_block = self.allocate()
        # 触发 GPU → GPU 内存拷贝 (nvMemcpy 或 custom kernel)
        copy_block_kernel(block_id, new_block, self.block_size)
        self.ref_counts[block_id] -= 1
        return new_block

    @property
    def available_blocks(self) -> int:
        return len(self.free_blocks)

4.3 内存节省的量化分析

以一个 7B 模型为例： - num_layers = 32, num_kv_heads = 32, head_dim = 128 - 每个 token 的 KV Cache 大小 = 2 × 32 × 32 × 128 × 2 bytes (fp16) = 524,288 bytes ≈ 0.5 MB - 若预留 8192 token 最大长度：8192 × 0.5 MB ≈ 4 GB 每请求 - 16 个并发请求：64 GB 浪费在不使用的预留空间

PagedAttention 将浪费从预留制改为按需制： - 实际使用的 token 才分配物理空间 - 内存利用率从 ~20-60% 提升到 ~95%+ - 相同 GPU 显存下并发吞吐量提升 2-3x

五、CUDA 实现：分页 Attention Kernel

5.1 传统 Attention vs Paged Attention

传统 Attention 内核直接从连续的 KV Cache tensor 中 gather 数据：

// 传统 attention：KV 内存连续
// V[seq_len, num_heads, head_dim] 是连续 tensor
__global__ void standard_attention_kernel(
    const half* Q,        // [num_heads, head_dim]
    const half* K,        // [seq_len, num_heads, head_dim] 连续
    const half* V,        // [seq_len, num_heads, head_dim] 连续
    half* output,
    int seq_len,
    int head_dim
) {
    int head_idx = blockIdx.x;
    int tid = threadIdx.x;

    float score = 0.0f;
    // 逐 token 计算 attention score + weighted sum
    for (int i = 0; i < seq_len; i++) {
        // 内存连续访问，合并访存友好
        float qk = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            qk += __half2float(Q[head_idx * head_dim + d]) * 
                  __half2float(K[i * num_heads * head_dim + head_idx * head_dim + d]);
        }
        score += expf(qk);
        // ... accumulate V
    }
}

PagedAttention 需要从分散的 block 中读取 KV 数据：

5.2 PagedAttention v2 Kernel

vLLM 使用自定义 CUDA kernel 实现分页访问，核心思路：

// PagedAttention：KV 分散在不同物理 block 中
__global__ void paged_attention_v2_kernel(
    const half* __restrict__ q,           // query [num_heads, head_dim]
    const half* __restrict__ k_buffer,    // 全局 KV 物理内存池 [num_blocks, block_size, num_heads, head_dim]
    const half* __restrict__ v_buffer,
    const int* __restrict__ block_tables,  // 块索引 [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ seq_lens,      // 每个序列的实际长度
    half* __restrict__ output,
    int num_heads,
    int head_dim,
    int block_size,
    int max_num_blocks_per_seq
) {
    // 每个 block 负责一个 head 的 attention 计算
    int head_idx = blockIdx.y;
    int seq_idx = blockIdx.x;
    int tid = threadIdx.x;

    int seq_len = seq_lens[seq_idx];
    int num_blocks = (seq_len + block_size - 1) / block_size;

    extern __shared__ float smem[];
    float* partial_max = smem;
    float* partial_sum = smem + blockDim.x;
    float* partial_out = smem + 2 * blockDim.x;

    float max_score = -INFINITY;
    float sum_exp = 0.0f;
    float acc[HEAD_DIM / 32] = {0};  // 假设 head_dim=128, 每次读 4 个 half2

    // Iterate over 每个物理 block
    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        int physical_block = block_tables[seq_idx * max_num_blocks_per_seq + block_idx];
        int token_offset = block_idx * block_size;

        // 每个线程处理 block 中的若干 token
        for (int tok = tid; tok < block_size && (token_offset + tok) < seq_len; tok += blockDim.x) {
            // 从物理 block 中读取 KV (非连续访存，但 block 内部连续)
            const half* k_ptr = k_buffer + 
                ((physical_block * block_size + tok) * num_heads + head_idx) * head_dim;
            const half* v_ptr = v_buffer + 
                ((physical_block * block_size + tok) * num_heads + head_idx) * head_dim;

            // 计算 Q·K^T
            float qk = 0.0f;
            #pragma unroll
            for (int d = 0; d < head_dim; d += 8) {
                half8 q8 = *reinterpret_cast<const half8*>(&q[head_idx * head_dim + d]);
                half8 k8 = *reinterpret_cast<const half8*>(&k_ptr[d]);
                qk += __half2float(q8.x) * __half2float(k8.x);
                qk += __half2float(q8.y) * __half2float(k8.y);
                // ... 展开 8 个 half 相乘
            }
            qk *= rsqrtf((float)head_dim);

            // Online softmax: 更新 max 和 sum
            float new_max = fmaxf(max_score, qk);
            float exp_diff = expf(max_score - new_max);
            float exp_qk = expf(qk - new_max);

            // 使用 warp 级别的协作归约
            sum_exp = sum_exp * exp_diff + exp_qk;
            max_score = new_max;

            // 加权累加 V（按 exp_diff 缩放）
            #pragma unroll
            for (int d = 0; d < head_dim; d++) {
                partial_out[d] = partial_out[d] * exp_diff + 
                    __half2float(v_ptr[d]) * exp_qk;
            }
        }
    }

    // 归一化 + 写回
    float inv_sum = 1.0f / sum_exp;
    #pragma unroll
    for (int d = 0; d < head_dim; d++) {
        output[seq_idx * num_heads * head_dim + head_idx * head_dim + d] = 
            __float2half(partial_out[d] * inv_sum);
    }
}

5.3 性能关键：Block Size 的选择

Block Size 是一个硬件相关的超参数： - 太小（< 8）：Block Table 变大，间接寻址开销增加；cache line 利用率下降 - 太大（> 64）：内部碎片增加；Copy-on-Write 的粒度变粗，共享效率下降

实践中最优值通常为 16-32 tokens。FlashAttention 的 block size（128-256 sequence positions）与 PagedAttention 不同——前者针对序列级并行，后者针对内存管理。

六、Prefix Caching：跨请求的语义级 Page Sharing

6.1 问题场景

在生产环境中，大量请求共享相同 System Prompt：

Request 1: "You are a helpful assistant. Translate to Chinese: Hello" → "你好"
Request 2: "You are a helpful assistant. Translate to Chinese: World" → "世界"
Request 3: "You are a helpful assistant. Translate to Chinese: Kernel" → "内核"

如果没有 Prefix Caching，每个请求都会重新计算 System Prompt 的 KV Cache（prefill 阶段），在 10K token system prompt 场景下这是巨大浪费。

6.2 基于哈希的自动前缀匹配

vLLM 的 Automatic Prefix Caching（APC）实现了基于内容哈希的 KV Cache 复用：

class PrefixCachingManager:
    def __init__(self, block_size: int = 16, hash_fn: str = "sha256"):
        self.block_size = block_size
        # 哈希缓存：block_hash → (物理块 ID, 访问时间戳)
        self.hash_to_block: Dict[str, Tuple[int, float]] = {}
        self.lru_heap: List[Tuple[float, int]] = []  # (last_access_time, block_id)

    def compute_block_hash(self, token_ids: Tuple[int], prev_block_hash: Optional[str]) -> str:
        """
        增量哈希：每个 block 的哈希包含前一个 block 的哈希
        这样相同 token 序列具有相同哈希值，且保护前缀语义
        """
        hasher = hashlib.sha256()
        hasher.update(struct.pack("i", len(token_ids)))
        for tid in token_ids:
            hasher.update(struct.pack("i", tid))
        if prev_block_hash:
            hasher.update(prev_block_hash.encode())
        return hasher.hexdigest()[:16]

    def match_prefix(self, token_ids: List[int]) -> Tuple[int, List[int]]:
        """返回匹配到的 token 数 和 已存在的物理 blocks"""
        matched_tokens = 0
        existing_blocks = []
        prev_hash = None

        for i in range(0, len(token_ids), self.block_size):
            block_tokens = tuple(token_ids[i:i + self.block_size])
            block_hash = self.compute_block_hash(block_tokens, prev_hash)

            if block_hash in self.hash_to_block:
                block_id, _ = self.hash_to_block[block_hash]
                # 类型断言确保 block_id 是 int
                assert isinstance(block_id, int)
                existing_blocks.append(block_id)
                matched_tokens += self.block_size
                prev_hash = block_hash
            else:
                break  # 前缀在此处中断

        return matched_tokens, existing_blocks

    def cache_blocks(self, token_ids: List[int], physical_blocks: List[int]):
        """将新计算的 KV blocks 注册到哈希缓存"""
        prev_hash = None
        for i, block_id in enumerate(physical_blocks):
            block_tokens = tuple(token_ids[i * self.block_size:(i + 1) * self.block_size])
            block_hash = self.compute_block_hash(block_tokens, prev_hash)
            self.hash_to_block[block_hash] = (block_id, time.time())
            prev_hash = block_hash

6.3 收益量化

Prefix Caching 的实际收益： - TTFT 降低：共享前缀部分的 prefill 时间完全消除。对于 8K system prompt 场景，TTFT 减少 40-70% - GPU 利用率提升：释放的 KV Cache 空间可接纳更多并发请求 - 请求级尾延迟改善：P99 latency 下降，因为 prefill 不再阻塞 decode

七、前沿演进：分离式架构与异构调度

7.1 Disaggregated Prefill-Decode

最新趋势（如 SGLang、DistServe 论文）将 Prefill 和 Decode 阶段分离到不同 GPU 上：

┌─────────────────┐         ┌─────────────────┐
│  Prefill GPU(s) │────────▶│  Decode GPU(s)  │
│  计算密集       │  KV     │  访存密集       │
│  大 batch       │ Transfer│  小 batch × N   │
│  A100/H100      │  RDMA   │  L40S/A10       │
└─────────────────┘         └─────────────────┘

原因：Prefill 是计算受限（大矩阵乘法），Decode 是带宽受限（逐 token 自回归）。两者混合调度时，一个长 prefill 会延迟所有 decode token 的生成。分离后各取所需：

• Prefill 用高端卡（H100）最大化计算吞吐
• Decode 用性价比卡（L40S）高并发低成本
• KV Transfer 通过 RDMA/NVLink 异步传输

7.2 Mooncake 月饼引擎：KV Cache 全局池化

Mooncake（月之暗面开源）将 KV Cache 存储层从 GPU 解耦为全局分布式存储：

                    ┌─────────────────┐
                    │   KV Cache Store │
                    │  CPU Mem + SSD   │
                    │  + GPU HBM 分层  │
                    └────────┬────────┘
                             │
            ┌────────────────┼────────────────┐
            ▼                ▼                ▼
     ┌────────────┐  ┌────────────┐  ┌────────────┐
     │ GPU Pool A │  │ GPU Pool B │  │ GPU Pool C │
     │ (Decode)   │  │ (Prefill)  │  │ (Idle→Act) │
     └────────────┘  └────────────┘  └────────────┘

八、工程实践与调优指南

8.1 关键参数配置

# vLLM 生产环境配置示例
model: meta-llama/Llama-2-70b-chat-hf
gpu_memory_utilization: 0.90   # KV Cache 占 HBM 比例
max_model_len: 4096            # 最大序列长度
max_num_seqs: 256              # 最大并发请求数
block_size: 16                 # PagedAttention block 大小
enable_prefix_caching: true    # 开启前缀缓存
enable_chunked_prefill: true   # 开启分块 prefill
max_num_batched_tokens: 2048   # 单次 iteration 最大 token 数

8.2 监控与告警

推理引擎需要监控的关键指标：

指标	说明	告警阈值
gpu_cache_usage_perc	KV Cache 利用率	> 90%
prefix_cache_hit_rate	前缀缓存命中率	< 30%
avg_time_to_first_token	平均 TTFT	> 200ms
avg_time_per_output_token	平均 decode 间隔	> 50ms
queue_wait_time	Scheduler 排队时间	> 100ms
num_preemptions	抢占次数/分钟	> 10

8.3 反模式与踩坑

反模式一：无脑开大 max_num_seqs

并发请求越多，每个请求可分配的 KV Cache 越少，抢占越频繁。一旦进入"抢占循环"（持续 swap ↔ restore），吞吐量反而崩溃。应该根据模型大小和预设 SLO 做容量规划。

反模式二：忽视 block_size 对 COW 效率的影响

Beam Search 场景下，多条 candidate 共享前缀的 KV blocks。如果 block_size 太小，意味着更多 block 需要 COW 分裂；如果太大，前缀浪费更多空间。7B 模型推荐 16，13B+ 推荐 32。

反模式三：Prefix Cache 的哈希冲突误用

自定义哈希函数时省略 prev_block_hash 会导致不同 token 序列产生相同哈希——"abc" + "def" 和 "abcd" + "ef" 可能被错误匹配。正确实现必须将前一个 block 的哈希作为当前输入。

九、结语

Continuous Batching 和 PagedAttention 不是什么神秘黑科技——它们是将操作系统 50 年成熟的虚拟内存管理思想，移植到 GPU 这一新型计算架构中的一次漂亮实践。

PagedAttention 的本质是分页内存管理。 Preemption 的本质是抢占式多任务调度。 Prefix Caching 的本质是内容寻址缓存。

理解这一点，不仅有助于调优推理引擎，更能让我们在面对未来新硬件（CXL 内存、异构计算、光子互连）时，更快找到适配的软件抽象层。毕竟，计算机科学的历史就是一个不断在旧概念中发现新应用的历史。