大模型推理加速与显存优化：PagedAttention 虚拟内存分页、KV Cache 压缩与连续批处理（Continuous Batching）实战

在大语言模型（LLM）从模型训练走向超大规模生产推理上线的过程中，GPU 显存（VRAM）与每秒处理 Token 数（Throughput）是决定云厂商与企业 AI 基础设施运营生死的核心指标。

许多刚接触 LLM 私有化部署（如 vLLM、SGLang、TGI）的工程师经常面临两大严峻现实：

1. 显存利用率极低但频繁 OOM：一张 80GB 的 A100/H800 显卡，加载完 70B 模型的权重（约 35GB FP8）后，剩下的 45GB 显存只能并发跑 不到 8 个并发请求，显存碎片率高达 60%~80%；
2. 队头阻塞与静态批处理低效（Head-of-Line Blocking）：传统的静态 Batching 模式下，一个 Batch 内的 4 个请求中如果有 1 个只需要生成 10 个 Token，而另 1 个需要生成 2,000 个 Token，短请求必须被迫在 GPU 里“陪跑死等”，造成算力严重浪费！

UC 伯克利团队提出的 PagedAttention（分页注意力机制） 与 连续批处理（Continuous Batching） 彻底重构了大模型推理引擎的底层基础设施，将 GPU 吞吐量直接暴拉 200% ~ 400%。

本文深入剖析自回归 KV-Cache 膨胀机理、操作系统的虚拟内存分页哲学在 GPU 显存上的迁移，并给出生产级 Python 连续批处理调度器与页表映射实战。

一、传统静态推理 vs PagedAttention 连续批处理对比矩阵

二、PagedAttention 虚拟分页内存映射底层机理

在传统的自回归计算中，每个 Token 生成时产生的 Key 和 Value 向量必须被保留供后续 Token 计算注意力。

PagedAttention 将操作系统分页哲学完美搬上了 GPU：

[逻辑上下文序列 (Logical Tokens)]
Token 0, 1, 2, ..., 15  | Token 16, 17, ..., 31 | Token 32, 33, ...
    (逻辑块 0)          |     (逻辑块 1)          |     (逻辑块 2)
         |              |          |              |          |
         v              |          v              |          v
+-------------------------------------------------------------------------------+
| 🌟 请求专属页表 (Block Table):                                                |
| - 逻辑块 0 ====> 物理显存块 Block #7                                         |
| - 逻辑块 1 ====> 物理显存块 Block #3                                         |
| - 逻辑块 2 ====> 物理显存块 Block #19                                        |
+-------------------------------------------------------------------------------+
                                 |
                                 v
+-------------------------------------------------------------------------------+
| GPU 显存非连续物理块池 (Physical VRAM Block Pool)                             |
| [Block 0] [Block 1] ... [Block #3] ... [Block #7] ... [Block #19]             |
+-------------------------------------------------------------------------------+

由于物理块不需要在显存中保持连续，推理引擎可以按需动态向 GPU 申请物理页，彻底消除了外部显存碎片化；在多轮对话分支或束搜索（Beam Search）中，不同的生成序列可以通过页表直接引用相同的物理 Block（零内存拷贝）！

三、生产级 Python 连续批处理（Continuous Batching）调度引擎实现

下面的 Python 实现构建了一套支持按需分配物理页、迭代级动态插拔请求、以及零队头阻塞的工业级连续批处理调度器。

"""
continuous_batching_scheduler.py
生产级大模型推理引擎核心调度器：PagedAttention 物理块分配与连续批处理实战
"""

import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple


BLOCK_SIZE = 16  # 每个物理页容纳 16 个 Tokens


@dataclass
class GenerationRequest:
    request_id: str
    prompt_tokens: int
    max_output_tokens: int
    generated_tokens: int = 0
    is_finished: bool = False
    physical_block_ids: List[int] = field(default_factory=list)


class PhysicalMemoryBlockPool:
    """GPU 显存物理块管理器"""

    def __init__(self, total_blocks: int = 100):
        self.total_blocks = total_blocks
        self.free_blocks: Set[int] = set(range(total_blocks))

    def allocate_block(self) -> Optional[int]:
        if not self.free_blocks:
            return None
        return self.free_blocks.pop()

    def free_block(self, block_id: int):
        self.free_blocks.add(block_id)

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


class ContinuousBatchingEngine:
    """连续批处理调度引擎中枢"""

    def __init__(self, total_vram_blocks: int = 64):
        self.memory_pool = PhysicalMemoryBlockPool(total_vram_blocks)
        self.waiting_queue: List[GenerationRequest] = []
        self.running_batch: List[GenerationRequest] = []

    def submit_request(self, req_id: str, prompt_len: int, max_tokens: int):
        req = GenerationRequest(
            request_id=req_id,
            prompt_tokens=prompt_len,
            max_output_tokens=max_tokens
        )
        self.waiting_queue.append(req)
        print(f"📥 [SUBMIT] 收到新请求: {req_id} (Prompt: {prompt_len}T, Target: {max_tokens}T)")

    def step_iteration(self):
        """执行单步自回归 Iteration 生成与动态批处理插拔"""
        # 1. 尝试从等待队列调度新请求加入 Running Batch (只要物理显存块充足)
        while self.waiting_queue:
            candidate = self.waiting_queue[0]
            # 计算该请求初始所需的物理块数: ceil(prompt_tokens / BLOCK_SIZE)
            needed_blocks = (candidate.prompt_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE

            if self.memory_pool.free_count >= needed_blocks:
                self.waiting_queue.pop(0)
                # 分配物理块
                for _ in range(needed_blocks):
                    candidate.physical_block_ids.append(self.memory_pool.allocate_block())
                self.running_batch.append(candidate)
                print(f"⚡ [JOIN BATCH] 请求 {candidate.request_id} 获得显存块，加入并发 Batch！")
            else:
                # 显存暂时不足，等待下一轮
                break

        if not self.running_batch:
            return

        print(f"\n⚙️ [ITERATION STEP] 当前并发执行 Batch 大小: 【{len(self.running_batch)} 个请求】 (可用显存块: {self.memory_pool.free_count})")

        # 2. 模拟 GPU 前向计算生成 1 个 Token
        finished_requests = []
        for req in self.running_batch:
            req.generated_tokens += 1
            total_current_tokens = req.prompt_tokens + req.generated_tokens

            # 检查是否需要申请新的物理块
            current_allocated_capacity = len(req.physical_block_ids) * BLOCK_SIZE
            if total_current_tokens > current_allocated_capacity:
                new_blk = self.memory_pool.allocate_block()
                if new_blk is not None:
                    req.physical_block_ids.append(new_blk)
                else:
                    print(f"🚨 [VRAM EXHAUSTED] 请求 {req.request_id} 遭遇显存不足，触发暂停调度！")

            # 检查是否生成完毕
            if req.generated_tokens >= req.max_output_tokens:
                req.is_finished = True
                finished_requests.append(req)

        # 3. 🌟 连续批处理核心: 完结请求立刻移出并释放显存物理块，零队头等待！
        for finished in finished_requests:
            self.running_batch.remove(finished)
            for blk_id in finished.physical_block_ids:
                self.memory_pool.free_block(blk_id)
            print(f"🎉 [FINISHED & FLUSH] 请求 {finished.request_id} 已生成完毕 ({finished.generated_tokens} Tokens)，立即释放 {len(finished.physical_block_ids)} 个物理块！")

生产演练与长短请求混合无阻塞调度展示

if __name__ == "__main__":
    engine = ContinuousBatchingEngine(total_vram_blocks=20)

    # 提交三个长短差异极大的请求
    # 请求 A: 短请求 (只需生成 2 个 Token)
    engine.submit_request("REQ_SHORT_A", prompt_len=10, max_tokens=2)
    # 请求 B: 长请求 (需要生成 5 个 Token)
    engine.submit_request("REQ_LONG_B", prompt_len=12, max_tokens=5)

    # 执行自回归调度模拟
    for step in range(1, 7):
        print(f"\n--- [自回归解码第 {step} 步] ---")
        engine.step_iteration()
        time.sleep(0.05)

        # 在第 3 步时动态插入一个新请求 C (模拟在线流式突发流量)
        if step == 2:
            engine.submit_request("REQ_NEW_C", prompt_len=8, max_tokens=3)

四、生产避坑与显存调优落地红线

在配置与运维生产大模型推理集群（vLLM / TensorRT-LLM）时，必须坚守以下四项实战准则：

1. 合理设置 gpu_memory_utilization（通常 0.85 ~ 0.90）：
   切勿将该参数配置为 1.0！必须预留 10%~15% 的显存用于 CUDA Context、PyTorch 激活值与临时 Buffer，防止推理突发 OOM Crash。
2. 选择适合业务的 Block Size（推荐 16 或 32）：
   • 块太小（如 4）：页表开销增加，GPU 内核寻址损耗大；
   • 块太大（如 128）：产生较多内部碎片。16 或 32 是当前工业界最高效的平衡点。
3. 开启 Chunked-Prefill 消除长 Prompt 对解码的抖动干扰：
   当一个 10k Token 的长 Prompt 到来时，传统的 Prefill 计算会霸占 GPU 算力数秒，导致其他正在生成的请求发生剧烈卡顿。开启 Chunked-Prefill（分片预填充） 能够将长 Prompt 打散在多个 Iteration 中平滑消化。

通过深入掌握 PagedAttention 的虚拟内存分页机理与迭代级连续批处理调度体系，AI 基础架构团队能够将昂贵的 GPU 显存与计算核心压榨到极致，实现千级别高并发、低延迟的大模型企业级生产服务交付。