什么是 vLLM V1？
2025 年 1 月，vLLM 团队宣布了 vLLM V1 的 alpha 版本：这是 vLLM 内部架构的一次重大重构。V1 的设计目标是：(a) 简化代码库，(b) 增强 vLLM 的可扩展性，以及 (c) 默认开启所有性能优化。最后一点对于编写优化注意力算子的内核开发者尤为重要。
vLLM 提供了广泛的性能优化：持续批处理 (continuous batching)、分页注意力 (paged attention)、投机采样 (speculative decoding)、分块预填充 (chunked prefill) 和前缀缓存 (prefix caching)。在 V0 中，其中一些优化方案相互兼容，而另一些则不兼容。要使它们协同工作，需要对代码库进行重大调整，并一直延伸到内核级别。
为了简化调度相关的代码库部分，V1 还对请求批处理的形成方式进行了重大更改。在 V0 中，vLLM 的调度器总是形成“预填充批次”（例如，仅由需要预填充的等待请求组成的批次）或“解码批次”（例如，仅由正在解码的运行中请求组成的批次）。而在 V1 中，调度器可以形成混合批次，其中可能包含各种状态的请求：新的预填充、分块预填充、解码，甚至是投机解码。
图 1：vLLM V1 中的批次与 vLLM V0 中批次的对比。
最初，唯一支持这些“V1 批次”的注意力后端是 FlashAttention 软件包 的 CUDA 版本。因此，V1 最初只能在 NVIDIA GPU 上使用，而不支持 AMD 等其他供应商的 GPU。vLLM V0 使用的 AMD GPU 注意力内核是作为自定义 C++/HIP 代码实现的，这使得它们很难适应 V1 的应用场景。正因如此，AMD、IBM Research 和 Red Hat 的团队决定为 vLLM V1 开发一套基于 Triton 内核的全新注意力后端，从而以一种平台可移植且对开发者更友好的方式实现对 AMD 的支持。
术语和基准
在深入了解内核工作原理之前，理解一些基本术语非常重要。当我们谈论 vLLM 中的序列时，有三个关键量。首先，上下文长度 (context length) 表示已经计算出注意力且 KV 张量已存储在分页 KV 缓存中的 token 数量。其次，查询长度 (query length) 描述了在此调度迭代中需要计算的“新”token 数量。最后，“序列长度”表示上下文长度与查询长度之和。
图 2：理解内核所需的一些基本术语
为了理解上述内容，思考一些边界情况很有帮助：
- 上下文长度 = 0，这对应于“纯”预填充操作。在这种情况下，KV 缓存中没有与该序列相关的内容，必须从头开始对提示词进行注意力计算。
- 上下文长度 > 0 且查询长度 ~ 1000 个 token。 这是典型的“分块预填充”操作，我们将长提示词拆分为较小的块。在这种情况下，新的 token 需要对已经存储在分页 KV 缓存中的旧 token 进行注意力计算。
- 上下文长度 > 0 且查询长度 = 1 个 token。 这是经典的“解码”操作。
- 上下文长度 > 0 且查询长度 ~ 3 个 token。 这是在使用小型草稿模型进行投机采样时出现的典型模式，该草稿模型为接下来的几个 token 生成估计值，并由主模型进行验证。
vLLM V1 的 Triton 内核需要支持所有这 4 种情况（以及介于两者之间的任何情况！）。
最初，唯一支持所有这些场景的 Triton 内核是 prefix_prefill 内核。该内核源自 LightLLM 项目，为在 AMD GPU 上运行 vLLM V1 提供了一种极佳的方式。与本文描述的所有注意力内核一样，它实现了分块 softmax 算法（具体为 FlashAttention-2 论文 中描述的算法）。在启动内核时，我们使用大小为 (batch_size, num_query_heads, query_length // BLOCK_M) 的网格，其中 BLOCK_M 是设置为 128 的常量。因此，前缀预填充将工作并行化分布在批次的序列、查询头和查询维度的“M 块”中。每个程序负责沿序列长度维度计算分块 softmax。本质上，它以“块”为单位遍历序列长度维度，其中序列维度的块大小由 BLOCK_N（设置为 64 的常量）定义，查询长度维度的块大小由 BLOCK_M 定义。对于对应于上下文的块，内核从分页 KV 缓存中读取 KV 张量。对于对应于查询的块，KV 张量从连续内存中读取（因为它们是在同一次前向传递中计算的）。图 3 展示了给定程序所执行的工作。
图 3：前缀预填充的一个程序所执行的工作（未显示批次和头维度的并行性）。
虽然该内核支持上述 4 种情况，但当 IBM Research、Red Hat 和 AMD 的团队开始运行我们的标准性能基准测试时，我们注意到 V1 的性能比基准 V0 性能慢了约 6 倍。
Triton 内核优化
AMD、IBM Research 和 Red Hat 的团队随后着手进行一系列优化，以尝试提高 vLLM V1 在 AMD GPU 上的性能。我们将这些优化分为 3 个阶段：
- 前缀预填充内核的优化。
- 使用专用解码序列内核的拆分方法。
- 用于预填充和解码（及介于两者之间）的全新统一 Triton 内核。
我们现在将讨论这些阶段。
阶段 1：前缀预填充内核的优化
此阶段的优化由 AMD 执行，旨在提高上述基准前缀预填充内核的性能。
将 Warp 数量从 8 减少到 4
在前缀预填充内核中，存在极高的溢出率 (spilling rate)，这严重影响了执行速度。AMD MI3xx GPU 只有 512 个向量通用寄存器 (VGPRs)，它们均匀分布在调度到特定 SIMD 的 Warp 上。如果内核需要更多寄存器（例如通过处理过大的内存块），一些寄存器就会被冲刷到暂存内存 HBM 中并在以后重新用于立即执行，稍后会恢复溢出的数据进行计算。8 个 Warp 意味着 2 个 Warp 调度在一个 SIMD 上，因此每个 Warp 有 256 个 VGPR。我们减少了活跃 Warp 的数量以消除寄存器溢出。仅此一项更改就带来了约 3-5 倍的加速，并且是大幅提升内核性能的最有效手段之一。
参数化调度和自动调优的内核配置
第一步减少了 SIMD 上 Warp 执行的并行性，以及 Warp 调度程序交错获取、解码和执行指令的可能性。为了增加灵活性，需要重构内核，以便正确扩展由一个程序处理的数据。我们公开了关键的内核启动参数（例如块大小、Warp 数量、流水线阶段、展开、预取等），以便自动调优器 (autotuner) 可以探索针对序列长度、KV 缓存页大小和硬件特性的最佳配置。通过自动调优，我们为许多现实场景找到了最合适的配置。自动调优本身在确定调用哪个内核以及它是否已经在缓存中时会增加延迟。为了避免这种情况，我们决定直接将自动调优的参数发送给内核。
缓存循环中对齐的块内存访问
分页 KV 缓存布局并不简单，它是 key_cache (num_blocks, num_heads, head_size // x, block_size, x) 和 v_cache (num_blocks, num_heads, head_size, block_size)。这种布局使得生成加载指令内部的偏移量并对其进行向量化变得困难，特别是对于 key 而言。加载向量化不仅对内核性能有重大贡献，还能降低寄存器压力。在处理先前计算的预填充块的内部缓存读取循环内，我们重构了加载逻辑以使用常量偏移量（即在编译时计算），以便 Triton 编译器能最好地了解要加载的数据。这迫使编译器生成向量化的全局加载，从而提高了内存带宽效率。此外，这种对齐解锁了编译器可以利用的循环展开和流水线机会，进一步减少了内存密集型区域的开销。
重构在线 Softmax 逻辑
在 AMD，我们注意到在线 Softmax 的实现方式并不理想。它对 P 尺度和 Acc 尺度进行了额外的计算。计算发生在热点循环内，导致了额外的寄存器使用和计算时间（向量操作是同步的）。我们重写了前缀预填充内核中的 Softmax 部分，以使用更少的中间寄存器（例如，重组累加、重新排序操作）。由此产生的较低寄存器压力降低了溢出风险，并提高了 Triton 生成代码中的指令级并行度。
简化内部循环执行
作为优化工作的一部分，我们消除了内部循环中处理非边界块时的冗余边界条件检查。通过根据块位置专门化循环体，非边界循环现在可以在没有防御越界访问的条件分支的情况下执行。这减少了控制流分歧，并使编译器能够生成更高效的线性代码，从而提高了指令吞吐量和整体内核性能——这对于大型、对齐良好的工作负载尤其有效。
阶段 2：针对解码序列的拆分方法
此阶段的优化工作始于 Red Hat 和 IBM Research 团队注意到前缀预填充内核用于解码序列（上述情况 3）的方式效率极低。对于这些序列，查询长度始终为 1 个 token。如果我们考虑块在查询长度维度上的布局方式，我们可以立即发现由于 BLOCK_M 设置为 128，我们需要屏蔽除一行之外的所有行，这意味着我们执行了实际所需工作量的 100 倍。
图 4：前缀预填充对于解码序列并不高效。
虽然减少解码序列的 BLOCK_M 似乎理所当然，但天下没有免费的午餐，因为 vLLM 将解码和预填充打包在同一个批次中，而 BLOCK_M 的值是编译时常量，因此对于批次内的所有序列必须相同。因此，为解码序列优化 BLOCK_M 会对同一批次中的其他序列产生不利影响。
基于上述观察，IBM Research 团队尝试的第一件事是修改前缀预填充内核，使得如果查询长度为 1，程序立即终止。因此，前缀预填充仅用于计算批次中查询长度 > 1 的序列。然后，我们启动第二个新编写的 Triton 内核 paged_attention_2d，如果查询长度 > 1，它会立即终止，而对于查询长度 = 1 的序列（例如解码），它以高度优化的方式计算分页注意力。“2D”名称指的是我们启动内核时使用的网格：它是一个形状为 (batch_size, num_query_heads) 的网格。请注意，前缀预填充使用 3D 网格，额外的维度对应于查询长度维度。由于解码不暴露这种额外的并行级别，因此不需要它（稍后将详细介绍 3D 内核）。
图 5：paged_attention_2D 的一个程序所执行的工作
图 5 展示了该内核执行的工作。每个程序都在单个查询头上操作（显示在 y 轴上），并沿上下文长度维度（显示在 x 轴上）执行分块 softmax。由于该内核仅处理分页 KV 缓存，我们将 BLOCK_N 设置为 16 个 token，以与 vLLM 中使用的默认页大小对齐。由于查询维度中没有冗余的“分块”，我们没有执行数量级上不必要的工作。将此内核与前缀预填充结合使用，使吞吐量提高了约 3.7 倍。
然而，该内核在处理使用分组查询注意力 (GQA) 的模型时表现不佳。原因在于它将所有查询头视为独立的，即使它们共享同一个 KV 头。为了克服这一限制，我们在查询头维度引入了额外的分块（如图 6 所示）。我们没有使用 (batch_size, num_query_heads) 形状的网格启动内核，而是使用 (batch_size, num_kv_heads) 形状的网格。每个程序为 QpKV 个查询头执行分块 softmax 算法，其中 QpKV = num_query_heads / num_kv_heads。通过将共享同一 KV 头的查询头捆绑在同一个块中，我们可以确保从 GPU VRAM 移动到计算核心的数据更少。此外，我们发现通过将 BLOCK_Q 向上取整为 16 的倍数，我们可以确保矩阵乘法映射到矩阵核心，从而带来额外的性能提升。总的来说，我们发现这些与 GQA 相关的优化带来了约 25% 的吞吐量提升。
图 6：带有 GQA 优化的 paged_attention_2D 的一个程序所执行的工作
阶段 3：统一注意力内核
虽然性能现在有了显著提高，但“拆分”方法仍有不足之处。特别是，将工作“拆分”到两个内核中（一个用于解码序列，另一个用于其他所有情况）有许多弊端。首先，我们需要维护两个复杂内核的代码。其次，在低延迟场景中，我们可能会受到基于 CPU 的启动开销的负面影响。最后，“拆分”解决方案没有考虑查询长度不完全为 1 但可能仍然很小的投机解码情况。在这种情况下，使用 prefix_prefill 将再次效率极低。
为了克服这些限制，IBM Research 团队开发了 unified_attention_2d Triton 内核（参见图 7）。该内核再次将启动网格的形状重新定义为 (total_num_q_blocks, num_kv_heads)，其中 total_num_q_blocks 是一个与整个批次中查询长度之和成比例的量。每个程序沿序列长度维度（图 7 中的 z 轴）执行分块 softmax 算法。分块沿查询头维度执行，使用大小恰好为 QpKV 的 BLOCK_Q（显示在图 7 的 y 轴上），以及沿“扁平化”批次维度（显示在图 7 的 x 轴上）执行，使用 BLOCK_M = 16/QpKV。通过这种方法，我们将跨查询头和查询维度中 token 的工作捆绑到相同的矩阵乘法中，这意味着在 QpKV < 16 的情况下，我们可以将更多工作打包到相同的张量核心指令中。注意，扁平化批次维度中的分块（“Q 块”）必须仔细定义，以便每个 Q 块仅包含来自批次中一个序列的 token（请注意图 7 中每个查询末尾的一些 Q 块是如何被屏蔽的）。此约束是必要的；否则，块将需要同时处理多个序列，这会极大地复杂化逻辑。
图 7：unified_attention_2d 内核的一个程序所执行的工作。
我们现在拥有一个单一的内核，可以用于 V1 批次，仅包含 239 行代码，使其更易于长期维护。然而，AMD 团队的实验表明，在某些场景下，Triton 内核的表现仍然不佳。特别是，AMD 团队发现当输出 token 数量非常长时，为 vLLM V0 编写的 C++/HIP 解码内核仍然能提供更好的性能。我们确定原因是统一的 Triton 内核没有利用沿上下文长度维度的任何并行性。为了解决这个问题，IBM Research 团队开发了上述内核的扩展版本 unified_attention_3d，遵循 Flash-Decoding 的思路，沿上下文长度维度拆分工作，从而以最终缩减步骤为代价创造了额外的并行性。
当头数量维度 (num_kv_heads) 或扁平化批次维度 (tot_num_q_blocks) 中没有足够的并行性来完全占用 GPU 时，最需要 3D 内核。如果批次相对较小，或者批次内的查询长度非常短（例如纯解码），通常会发生这种情况。然而，如果批次非常大，或者批次包含一些非常长的查询，那么 2D 内核已经允许足够的并行性，执行最终缩减步骤的开销可能不值得。因此，IBM Research 团队实现了一个 简单的启发式方法 来决定何时使用 2D 内核与 3D 内核。
随后，AMD 团队在统一的注意力 2D/3D 内核之上贡献了额外的优化，使 vLLM V1 的性能达到了今天的水平。
SWA 循环边界调整和配置更新
优化了滑动窗口注意力 (SWA) 机制，严格将计算限制在活跃窗口范围内。此前，内核会评估整个序列长度的注意力分数，随后丢弃目标窗口之外的值，导致不必要的内存读取和计算开销。新实现引入了范围感知索引和掩码逻辑，动态限制查询-键交互到相关的 token 子集。这一变化不仅减少了长序列场景中的二次计算成本，还提高了内存局部性、缓存效率和整体吞吐量，特别有利于流式或增量解码工作负载。
KV 和 Q 块的 ‘cg’ 缓存修饰符（条件性重用）：
应用了缓存修饰符来优化全局内存加载——这在内存受限情形下的 AMD GPU 上特别有效。
用于改善 XCD 映射的 2D 注意力网格重排序：
重新设计了 2D 注意力网格布局以增强空间局部性和跨芯片数据重用。新的映射策略对块索引进行重新排序，使得处理相邻查询-键块的工作组优先共位于同一个 XCD（跨计算芯片）上。这种对齐确保共享重叠键/值区域的块在同一内存域上执行，显著改善了 L2 缓存驻留并减少了跨芯片流量。通过最小化远程内存访问并提高共享数据的时序重用，此优化增强了有效内存带宽，并在多芯片配置中带来了更一致的延迟缩放。
增加 BLOCK_M 以减少数据重新加载：
调整了块高度 (BLOCK_M) 以更好地平衡寄存器利用率和内存访问效率。通过增加每个块处理的行数，每个线程组现在可以在驱逐前重用加载的查询和键数据，从而显著减少冗余的全局内存读取。这一变化提高了计算与内存的比率，增强了算术强度，并提高了整体吞吐量——尤其是在内存带宽往往成为主要瓶颈的长序列长度情况下。新配置还允许内存预取和计算之间更好的重叠，进一步提升了硬件利用率。
专门的预填充配置调优：
为拆分内核执行模型中的预填充阶段工作负载引入了专门的参数预设。这些预设自动调整关键的内核参数——例如块尺寸、Warp 数量和流水线深度——以最大化针对预填充特性的 GPU 占用率和数据重用。通过将块大小和内存预取模式与预填充操作中通常较大的批次和序列维度对齐，新配置提高了缓存利用率，减少了全局内存流量，并保持了高算术强度。
修订后的内核选择逻辑（2D vs. 3D）：
增强了 vLLM 的内核选择启发式逻辑，使其能够根据当前的序列长度、批次形状和 token 分布更准确地在 2D 和 3D 启动配置之间进行选择。更新后的逻辑考虑了工作负载几何结构和 GPU 占用特性，确保每种启动策略都在其表现最佳的情形下应用。这种自适应选择改善了内核在各种输入模式下的效率——最大限度地减少了空闲线程，提高了内存访问规律性，并在小批次和大上下文工作负载中提供了更一致的性能。
扩展 3D 解码的网格拆分：
增加了沿第三个启动维度的划分程度，以生成更大的整体网格，从而提高 GPU 占用率和并行性。此调整允许同时调度更多的线程块，确保更好地利用可用的 SM，特别是在长序列解码工作负载中，块计算量可能会限制并发性。这一变化增强了设备上的负载均衡，并带来了随序列长度更一致的性能扩展。
性能基准测试
本博客描述了数月来贡献给开源社区的大量优化。每项优化的性能提升可以在上述文本链接的 PR 中找到。在本节中，我们旨在评估这些工作的累积效果，并展示 vLLM V1（基于 Triton）现在在 AMD GPU 上的真实异构服务基准测试中优于 V0（基于 C++/HIP）。
我们将比较 vLLM V0 和 V1 在单张 AMD MI300x GPU 上针对 mistralai/Mistral-Small-24B-Instruct-2501 模型的性能。所有实验均使用公开可用的 Docker 镜像和 vLLM 内置的基准测试工具执行，因此应该易于复现。我们分别使用了以下 V0 和 V1 的 Docker 镜像：
比较 V0 和 V1 的性能很复杂，因为两个版本有不同的默认设置。特别是，控制 vLLM 并发处理的最大序列数量的参数 max-num-seqs 在 V0 中默认设置为 256，而在 V1 中默认为 1024。同样，V0 中默认禁用分块预填充，而 V1 默认启用，且对于此特定模型，分块预填充参数 max-num-batched-tokens 默认值为 8192。这些参数可能会显著影响性能，并可能掩盖 V0 和 V1 之间与不同内核实现相关的差异。因此，我们运行了 4 个不同的实验来尝试控制这些影响，并隔离本文所述内核工作带来的收益。
V0（默认设置，禁用前缀缓存）
VLLM_USE_V1=0 vllm serve mistralai/Mistral-Small-24B-Instruct-2501 \ --no-enable-prefix-caching
V0（增加最大序列数）
VLLM_USE_V1=0 vllm serve mistralai/Mistral-Small-24B-Instruct-2501 \ --no-enable-prefix-caching \ --max-num-seqs 1024
V0（增加最大序列数并启用分块预填充）
VLLM_USE_V1=0 vllm serve mistralai/Mistral-Small-24B-Instruct-2501 \ --no-enable-prefix-caching \ --max-num-seqs 1024 \ --enable-chunked-prefill \ --max-num-batched-tokens 8192
V1（默认设置，禁用前缀缓存）
VLLM_USE_V1=1 vllm serve mistralai/Mistral-Small-24B-Instruct-2501 \ --no-enable-prefix-caching
客户端（基准测试）
在所有 4 个实验中，我们使用 vLLM 的内置基准测试工具以异构服务工作负载向服务器发起请求。
vllm bench serve \ --model mistralai/Mistral-Small-24B-Instruct-2501 \ --dataset-name sharegpt \ --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \ --ignore_eos
结果
在图 8、9 和 10 中，我们比较了上述 4 种配置的总 token 吞吐量、首字延迟 (TTFT) 和字间延迟 (ITL)。即使我们将 V0 的默认参数调整为与 V1 匹配，在使用为 vLLM V1 开发和优化的 Triton 内核时，我们在所有指标上仍看到了约 10% 的提升。
图 8：不同 vLLM 配置的总 token 吞吐量（token/秒）对比
图 9：不同 vLLM 配置的首字延迟 (TTFT)（秒）对比
图 10：不同 vLLM 配置的字间延迟 (ITL)（毫秒）对比。
结论
在本博客中，我们描述了 AMD、IBM Research 和 Red Hat 的团队如何通过编写优化的 Triton 内核并将其贡献给开源社区，为 vLLM V1 构建优化的注意力后端。我们已经展示了基于 Triton 的 vLLM V1 在 AMD MI300x GPU 上比使用自定义 C++/HIP 实现的 vLLM V0 可提供 10% 的吞吐量提升。我们希望本博客中提供的详细信息和见解对编写注意力或其他算子内核的 Triton 开发者有所帮助。如果您对如何进一步提高 vLLM 性能有任何想法，请随时联系我们！
鸣谢
这项工作由 3 个不同组织的庞大团队完成——感谢所有参与者。
IBM Research
Burkhard Ringlein, Jan van Lunteren, Chih-Chieh Yang, Sara Kokkila Schumacher, Thomas Parnell, Mudhakar Srivatsa, Raghu Ganti
AMD
Aleksandr Malyshev, Vinayak Gokhale, Mehmet Kaymak, Ali Zaidy, Joe Shajrawi
Red Hat
Sage Moore, Tyler Michael Smith, Robert Shaw, Lucas Wilkinson