PagedAttention:分页注意力——推理引擎的显存革命

一句话理解

PagedAttention 是vLLM的核心技术——借鉴操作系统的分页内存管理思想,用"按需分配"的方式管理KV Cache,让GPU显存利用率飙升到90%以上。

传统方式(预分配):
│████████████│                    │████████████│
│ 预分配1G    │                    │  浪费800M   │
└────────────┘                    └────────────┘

PagedAttention(动态分配):
│█│██│█│███│█│                   │█│██│█│███│█│
│按需分配     │                   │块式管理     │
│无浪费       │                   │灵活拼接     │

背景:为什么需要PagedAttention?

KV Cache的显存困境

回顾W19的内容:KV Cache是大模型推理的关键,但显存占用巨大:

以LLaMA-7B为例,处理一个请求:

- 32层Transformer
- 隐藏维度4096
- 半精度(FP16) = 2字节
- 上下文长度8192 tokens

显存 = 32层 × 4096 × 2字节 × 8192 × 2(K+V)
     ≈ 4GB显存

如果同时处理10个请求:
10 × 4GB = 40GB显存

A100才80GB显存,光KV Cache就用了一半!

传统方案的浪费

问题 说明
核心问题 预分配 + 固定分块,每个请求分配一个固定大小的"块",但实际需要的空间不同
请求A 生成了2000 tokens → 需要2000的KV空间,但预分配4000 → 浪费2000
请求B 生成了500 tokens → 需要500的KV空间,但预分配4000 → 浪费3500

总浪费:5500 tokens的显存!


PagedAttention的核心思想

操作系统分页的启发

┌────────────────────────────────────────────────────────┐
│                                                          │
│   操作系统如何管理内存?                                 │
│                                                          │
│   物理内存:    │▓│▓│▓│░│▓│░│░│▓│░│▓│                   │
│                 1  2  3  4  5  6  7  8  9  10            │
│                                                          │
│   虚拟内存:    │░░░░░░░░░░░░░░░░░░░░│                   │
│                 1                 1024                  │
│                                                          │
│   分页机制:虚拟页面 → 物理页面(不连续但逻辑连续)      │
│                                                          │
│   好处:                                                │
│   ├─ 不需要连续的物理内存                               │
│   ├─ 按需分配,不浪费                                   │
│   └─ 共享内存页,减少冗余                               │
│                                                          │
└────────────────────────────────────────────────────────┘

PagedAttention的类比

PagedAttention = 操作系统分页应用到KV Cache

特点 说明
KV Cache分页 把KV Cache分成固定大小的"块"(Block)
连续存储 每个请求的KV不需要连续存储
分配方式 按需分配,动态拼接
方案 图示 效果
传统方案 ████████████░░░░░░░░░ 浪费50%
PagedAttention ██│███│█│████│ 无浪费

拼接来自不同物理位置的块


技术原理

Block结构

# PagedAttention的Block抽象

class PagedAttentionBlock:
    """
    KV Cache的Block,类似操作系统的内存页
    """
    def __init__(self, block_size=16):
        self.block_size = block_size  # 每个Block容纳16个Token
        self.num_tokens = 0          # 当前Block中的Token数
        self.k_cache = torch.zeros(block_size, num_heads, head_dim)  # KV存储
        self.v_cache = torch.zeros(block_size, num_heads, head_dim)
        self.block_id = None         # 物理Block编号
        self.ref_count = 0           # 引用计数(用于共享)

动态分配

class KVCacheManager:
    """
    类似操作系统的虚拟内存管理器
    """
    def __init__(self, total_num_blocks=1000):
        # 物理Block池
        self.physical_blocks = [PagedAttentionBlock() 
                                for _ in range(total_num_blocks)]
        
        # 虚拟Block → 物理Block的映射
        self.block_mapping = {}
        
        # 空闲Block队列
        self.free_blocks = set(range(total_num_blocks))
    
    def allocate(self, num_tokens_needed):
        """
        动态分配KV Cache空间
        """
        # 计算需要多少Block
        num_blocks_needed = (num_tokens_needed + self.block_size - 1) // self.block_size
        
        # 分配空闲Block
        allocated_blocks = []
        for _ in range(num_blocks_needed):
            if not self.free_blocks:
                # Evict最近最少使用的Block
                evicted = self.evict_lru_block()
                allocated_blocks.append(evicted)
            else:
                block_id = self.free_blocks.pop()
                allocated_blocks.append(block_id)
        
        # 返回逻辑连续的虚拟地址(映射到物理Block)
        return VirtualBlock(allocated_blocks)
    
    def write(self, virtual_block, token_ids, k_values, v_values):
        """
        写入KV Cache
        """
        offset = 0
        for i, block_id in enumerate(virtual_block.block_ids):
            physical_block = self.physical_blocks[block_id]
            
            # 计算这个Block写入多少Token
            num_to_write = min(self.block_size, len(token_ids) - offset)
            
            # 写入KV
            physical_block.k_cache[:num_to_write] = k_values[offset:offset+num_to_write]
            physical_block.v_cache[:num_to_write] = v_values[offset:offset+num_to_write]
            
            physical_block.num_tokens = num_to_write
            offset += num_to_write

并行Attention计算

def paged_attention(query, block_mapping, kv_cache):
    """
    PagedAttention的Attention计算
    """
    output = torch.zeros_like(query)
    
    for block_id in range(num_blocks):
        # 动态获取物理Block
        physical_block = kv_cache.physical_blocks[block_id]
        
        # 只计算这个Block中有效的Token
        valid_len = physical_block.num_tokens
        
        # 计算注意力
        scores = query @ physical_block.k_cache[:valid_len].T / math.sqrt(d)
        attn_weights = F.softmax(scores, dim=-1)
        
        output += attn_weights @ physical_block.v_cache[:valid_len]
    
    return output

vLLM:PagedAttention的生产级实现

架构概览

组件 说明
Scheduler(调度器) 请求队列管理、Block分配/回收、Batch调度
PagedAttention Engine 动态Block管理、FlashAttention融合、CUDA Kernel优化

使用示例

from vllm import LLM, SamplingParams

# 初始化vLLM,自动使用PagedAttention
llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    
    # GPU显存利用配置
    gpu_memory_utilization=0.9,  # 90%用于KV Cache
    
    # 最大上下文长度
    max_model_len=32768,
    
    # Block大小(默认16)
    block_size=16,
    
    # 并发请求数
    max_num_seqs=256,
)

# 批量推理
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=512,
)

prompts = [
    "解释什么是PagedAttention",
    "写一个Python快速排序",
    "比较CNN和RNN的区别",
]

# vLLM会自动:
# 1. 批量处理请求
# 2. 动态管理KV Cache
# 3. 最大化GPU利用率
outputs = llm.generate(prompts, sampling_params)

性能对比

显存利用率

测试环境:LLaMA-7B on A100 80GB

方案 显存利用率 提升
HuggingFace (无PagedAttention) ~40% 基准
vLLM (PagedAttention) ~90% ↑ 125%

吞吐量对比

配置 HuggingFace vLLM 提升
7B模型 50 req/s 200 req/s 4倍
13B模型 25 req/s 100 req/s 4倍
70B模型 8 req/s 30 req/s 3.75倍

延迟对比

生成1000 tokens的平均延迟:

模型大小    HF延迟    vLLM延迟    节省
─────────────────────────────────
7B          2.0s      0.8s       60%
13B         4.0s      1.2s       70%
70B         12.0s     3.5s       71%

2026年最新进展

vLLM 0.6.x系列(2026年)

新特性:
1. 投机解码原生支持
   └─ PagedAttention + Speculative Decoding联合优化
   
2. Chunked Prefill
   └─ 将长prompt分块处理,避免显存峰值
   
3. 动态Block大小
   └─ 不同请求自动选择最优Block大小(8/16/32)
   
4. 多模态支持
   └─ 图像Token和文本Token统一管理

ChunkKV(长上下文优化)

2026年新发布的KV Cache压缩技术:

原理:
1. 将KV Cache按语义块分割
2. 只保留关键块的完整KV
3. 其他块用低秩近似压缩

效果:
- 显存占用减少50%
- 长上下文处理能力翻倍
- 精度损失<2%

ThinKV(推理模型优化)

专门为Reasoning Model设计的KV优化:

推理模型特点:
- CoT过程很长(可能几百个Token)
- 需要保留完整的推理历史

ThinKV方案:
- 检测关键决策点
- 压缩中间推理步骤
- 保持推理连贯性

技术对比

分页方案横向对比

技术 开发者 特点 适用场景
PagedAttention vLLM 固定Block,动态映射 通用推理
ChunkKV 学术 语义分块 长上下文
ThinKV 学术 推理链压缩 Reasoning Model
R-KV NVIDIA 递归压缩 超长序列

Block大小选择

Block越大:
├─ 优点:管理开销小,碎片少
└─ 缺点:内部碎片多

Block越小:
├─ 优点:碎片少,利用率高
└─ 缺点:管理开销大

推荐选择:
- 短序列(<4K):block_size=16
- 中等序列(4K-32K):block_size=16
- 长序列(>32K):block_size=32

实际应用

搭建高并发AI服务

# 使用vLLM搭建每秒处理1000+请求的服务

from vllm import LLLMEngine, SamplingParams
from fastapi import FastAPI
import uvicorn

app = FastAPI()

# 初始化引擎
engine = LLMEngine(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,  # 4卡并行
    gpu_memory_utilization=0.95,
    max_num_seqs=1000,
)

@app.post("/generate")
async def generate(request: dict):
    outputs = engine.generate(
        [request["prompt"]],
        SamplingParams(
            temperature=request.get("temperature", 0.7),
            max_tokens=request.get("max_tokens", 512),
        )
    )
    return {"text": outputs[0].outputs[0].text}

# 启动服务
# uvicorn app:app --host 0.0.0.0 --port 8000

长上下文问答

# 处理10万Token的文档

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    max_model_len=131072,  # 128K上下文
    gpu_memory_utilization=0.9,
    block_size=16,
)

# 自动处理长文档,不爆显存
response = llm.generate(
    "请总结这份文档的主要内容",
    SamplingParams(max_tokens=1024)
)

常见问题

Q1:PagedAttention和FlashAttention是什么关系?

互补关系

技术 解决的问题 层级
FlashAttention Attention计算的IO优化 计算层
PagedAttention KV Cache的显存管理 存储层
可以同时使用:
FlashAttention → 加速Attention计算
PagedAttention → 管理KV Cache存储

两者结合 = 最佳性能

Q2:使用vLLM需要注意什么?

# 1. 显存容量估算
estimated_memory = (
    model_size_gb +                    # 模型参数
    kv_cache_per_token_gb *            # 每个Token的KV
    max_context_len *                 # 最大上下文
    num_requests                      # 并发数
)

# 2. 动态调整
if oom_error:
    llm = LLM(
        gpu_memory_utilization=0.8,  # 降低显存使用
        block_size=32,                # 增大Block减少开销
    )

# 3. 监控指标
print(f"""
显存使用: {vllm_profile.gpu_memory_utilized:.1f}GB / {vllm_profile.gpu_memory_total:.1f}GB
Block数: {vllm_profile.num_used_blocks} / {vllm_profile.num_total_blocks}
前缀缓存命中率: {vllm_profile.prefix_cache_hit_rate:.1%}
""")

Q3:什么时候不用vLLM?

场景 原因 替代方案
单请求、低延迟 vLLM启动开销较大 HuggingFace
开发调试 vLLM日志不够友好 直接用HF
模型微调 vLLM只支持推理 DeepSpeed
小模型 收益不明显 直接用HF

总结

PagedAttention的核心价值:

  • 借鉴OS分页:按需分配KV Cache
  • 显存利用率:40% → 90%(提升125%)
  • 吞吐量提升:3-4倍
  • 延迟降低:60-70%
  • vLLM生产级实现:业界标配

技术栈:

  • PagedAttention = 存储层优化
  • FlashAttention = 计算层优化
  • Speculative Decoding = 生成加速
  • 三者结合 = 极致推理性能

2026年进化:

  • ChunkKV:长上下文压缩
  • ThinKV:推理模型优化
  • R-KV:NVIDIA原生加速

“没有PagedAttention,就没有AI的实时体验!”


延伸阅读

相关文章 说明
W19 KV Cache KV Cache基础
W20 Speculative Decoding 推理加速技术
W13 Transformer 注意力机制

本文收录于「AI词汇专栏」
相关阅读:W19 KV Cache · W20 Speculative Decoding

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐