vLLM中FlashAttention与KVCache交换机制解析

vLLM 中 FlashAttention 与 KVCache 交换机制深度解析

在当前大模型推理部署的工程实践中，高吞吐、低延迟、内存高效已成为衡量系统性能的核心指标。随着 LLM 应用从实验走向生产，我们不再满足于“能跑”，而是追求“跑得快、省资源、撑得住”。vLLM 正是在这一背景下脱颖而出的高性能推理框架——它不仅让 LLaMA-65B 这样的庞然大物能在单卡上实现每秒数十 token 的输出，更将显存利用率提升到接近理论极限。

这背后的关键，正是FlashAttention的极致计算优化，以及基于PagedAttention的 KVCache 动态管理机制。而连接这两者的“神经中枢”之一，就是那个看似简单却至关重要的swap_blocks操作。

要理解 vLLM 的设计哲学，不妨先回到最根本的问题：Transformer 推理到底“卡”在哪里？

标准注意力公式大家都很熟悉：

$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
$$

但当你把它放到 GPU 上跑长序列时，问题就来了。传统实现会把 $ QK^T $ 算出来并暂存在显存中，这个中间结果是 $ (N, N) $ 的矩阵。当序列长度达到 8k，仅这一项就要占用超过 250MB 显存（FP16），而且后续 Softmax 和乘 $ V $ 的过程还得反复读写这块数据。GPU 的算力越来越强，但显存带宽增长缓慢，“等数据”的时间远超“算数据”的时间——这就是典型的memory-bound场景。

更糟的是，Softmax 需要全局最大值和指数和，传统做法得先扫一遍求 max，再扫一遍算 sum，最后归一化，三次全局访存不说，还容易因为数值过大导致 float overflow。这些问题加在一起，使得传统 attention 成为推理效率的瓶颈。

于是 FlashAttention 出现了。它的核心思想非常直接：不让中间结果落地，把整个 attention 计算塞进一个 CUDA kernel 里完成，并通过分块策略减少对 HBM 的访问次数。

具体怎么做？三个关键词：Kernel Fusion + Tiling + OnlineSoftmax。

首先，不再分步执行 $ QK^T $ → Softmax → $ PV $，而是合并成一个 kernel，所有计算都在 SRAM（共享内存）中完成，避免频繁出入 HBM。

其次，采用 Tiling 技术，将 K 和 V 按 block 切片加载进 SRAM。比如每次只加载 256 个历史 token 的 Key/Value 向量，与当前 Query 做局部 attention，然后增量更新最终结果。

最关键的是OnlineSoftmax机制。它允许我们在不知道全局信息的前提下，边加载 block 边动态维护 softmax 所需的状态。假设有前缀的最大值 $ M_{\text{old}} $ 和归一化因子 $ L_{\text{old}} $，新来一块输入 $ X_i $，那么可以递推地更新：

$$
M_{\text{new}} = \max(M_{\text{old}}, \max(X_i)) \
L_{\text{new}} = L_{\text{old}} \cdot \exp(M_{\text{old}} - M_{\text{new}}) + \sum \exp(X_i - M_{\text{new}})
$$

同时输出也相应累加：

$$
O_{\text{new}} = \frac{L_{\text{old}} \cdot \exp(M_{\text{old}} - M_{\text{new}})}{L_{\text{new}}} O_{\text{old}} + \frac{\sum \exp(X_i - M_{\text{new}}) V_i}{L_{\text{new}}}
$$

这套机制实现了真正的 one-pass 流式计算，彻底摆脱了两遍扫描的束缚。这也是为什么 FlashAttention 能在保持数值稳定的同时，将 attention 从 memory-bound 推向 compute-bound。

不过，就算算得再快，如果存不下，依然白搭。

在服务端场景下，不同请求的上下文长度差异极大：有的用户可能只问一句“你好”，有的则上传整篇 PDF 提问。传统 KVCache 管理方式通常为每个 sequence 预分配一段连续显存空间。这种静态分配导致严重的内部碎片化——比如预分配 4k 空间，实际只用了 1.2k，剩下近 3k 就浪费了。尤其在高并发下，这种碎片累积起来足以让系统提前耗尽显存。

vLLM 的解决方案借鉴了操作系统的页式内存管理思想：引入PagedAttention，将每个 sequence 的 KVCache 拆分为多个固定大小的 “page”，每个 page 存储若干 token 的 K/V 向量，通过类似页表（page table）的结构进行逻辑索引。

这样一来，即使物理上分散存放，也能按需拼接使用。就像操作系统用虚拟内存突破物理内存限制一样，PagedAttention 让 vLLM 实现了对显存的“虚拟化”管理，显著提升了利用率。

但新的挑战随之而来：既然 pages 是动态分配的，那如何在运行时灵活迁移、重组这些 blocks？特别是在连续批处理（Continuous Batching）中，某些请求完成生成后，其占用的 pages 必须被快速回收并重新分配给新请求。

这就引出了关键操作：KVCache Swap，也就是swap_blocks。

我们可以看看这个函数的大致签名：

@staticmethod def swap_blocks( src_kv_cache: torch.Tensor, dst_kv_cache: torch.Tensor, src_to_dst: torch.Tensor, ) -> None:

参数含义如下：
-src_kv_cache: 源 KV 缓存池，通常是 GPU 上的主缓存区；
-dst_kv_cache: 目标缓存区，可能是另一个设备或子集；
-src_to_dst: 映射关系张量，src_to_dst[i] = j表示将第 i 个源 block 复制到第 j 个目标 block。

该函数的核心功能是：根据映射表批量复制 KV blocks。典型应用场景包括：

• 请求结束时，将其使用的 pages 标记为空闲，供后续复用；
• 新请求接入时，从 free pool 分配空闲 blocks 并建立索引；
• 内存整理（defragmentation）过程中，将离散的 pages 重新排列以形成更紧凑布局；
• 多 GPU 场景下同步分片状态。

你可能会问：Key 和 Value 为什么要分开交换？

虽然它们属于同一 attention 层，但在存储结构上通常是分离的张量（shape 或 layout 不同），且 FlashAttention 内核分别读取 K 和 V。更重要的是，这种设计保留了灵活性。例如在 MQA/GQA 架构中，多个 Query Head 共享一组 Key/Value Head，此时完全可以跳过多余的 Value copy，只做必要的索引映射即可。

为了更直观理解swap_blocks的行为，下面是一个简化版 CPU 实现：

import torch def mock_swap_blocks(src_cache: torch.Tensor, dst_cache: torch.Tensor, src_to_dst: torch.LongTensor): """ 模拟 swap_blocks 行为：将 src_cache 中的 block 按照 src_to_dst 映射复制到 dst_cache src_cache: [num_src_blocks, *block_shape] dst_cache: [num_dst_blocks, *block_shape] src_to_dst: [batch_size, max_blocks_per_seq], 值为 dst index，-1 表示无效 """ for dst_idx in range(dst_cache.size(0)): matching_src = (src_to_dst == dst_idx).nonzero(as_tuple=True)[0] if len(matching_src) == 0: continue src_idx = matching_src[0].item() if src_idx < src_cache.size(0): dst_cache[dst_idx] = src_cache[src_idx] # 示例 src_kv = torch.randn(8, 32, 64, 16) # 8 个 key blocks dst_kv = torch.zeros(4, 32, 64, 16) mapping = torch.tensor([2, -1, 5, 0]) # dst[0] <- src[2], dst[2] <- src[5], dst[3] <- src[0] mock_swap_blocks(src_kv, dst_kv, mapping) print("Swap completed. Target buffer filled:", (dst_kv.sum(dim=[1,2,3]) != 0).sum().item(), "blocks")

注：真实实现由 CUDA kernel 完成，直接操作显存地址，效率极高。

整个推理流程因此变得极为高效：

graph TD A[新请求到达] --> B{是否可加入当前Batch?} B -->|是| C[分配空闲KVCache Pages] B -->|否| D[启动新Batch] C --> E[Prefill: 使用FlashAttention计算全部KV并写入Pages] E --> F[Decode: 自回归生成Token] F --> G{是否有请求完成?} G -->|是| H[回收其Pages至Free Pool] G -->|否| F H --> I[分配给新请求] I --> C

在这个闭环中：
- FlashAttention 加速每一次 attention 计算；
- PagedAttention 实现细粒度的显存管理；
-swap_blocks支持运行时动态调度与复用；

三者协同作用，使 vLLM 在相同硬件条件下支持的并发请求数大幅提升，实测吞吐相比 HuggingFace Transformers 提升5~10 倍，部分场景甚至更高。

这也解释了为何越来越多的企业选择 vLLM 作为生产环境的推理引擎。它不只是“更快一点”，而是从根本上重构了 LLM 推理的资源模型。

从面试角度看，这类底层机制也常被深入考察：
- 为什么普通 softmax 无法实现 one-pass？因为它依赖全局信息，必须两次扫描；而 FlashAttention 借助 OnlineSoftmax 的递推性质，做到了真正的一次遍历。
- 如何优化 MQA/GQA？不能简单复制 KV，那样会造成冗余传输和存储。正确做法是在 kernel 层通过 indexing 直接寻址共享 head，结合 PagedAttention 可进一步节省显存。
-swap_blocks能否省略？绝不能。它是实现动态内存管理和高密度 batching 的基石。没有它，就只能退回到固定分配的老路，丧失灵活性。

可以说，vLLM 的成功并非偶然，而是对“算”与“存”两大维度深刻洞察的结果：

前者解决了“算得慢”，后者解决了“存不下”，而swap_blocks正是打通动态调度的关键环节。

对于希望部署 LLaMA、Qwen、ChatGLM 等主流模型的团队而言，vLLM 不仅提供了开箱即用的高性能能力，还通过 OpenAI 兼容 API 和量化支持（GPTQ/AWQ），实现了低成本、易集成、高可用的生产级解决方案。

掌握这些机制的意义，远不止于调优推理服务。它让我们看到，在大模型时代，系统级创新同样重要——算法的进步需要与底层架构的演进相匹配，才能真正释放潜力。