新媒网 · 跨境数据社区

2025年1月，vLLM团队宣布推出vLLM V1的Alpha版本，这标志着vLLM内部架构的一次重大重新设计。此次V1版本的设计目标明确而深远：一是简化代码库，提升可维护性；二是增强vLLM的扩展性，以适应更广泛的应用场景；三是默认开启所有性能优化，确保用户在开箱即用时即可享受最佳表现。对于专注于编写优化注意力内核的开发者而言，第三点尤为重要。

vLLM提供了多项先进的性能优化技术，包括连续批处理（continuous batching）、分页注意力（paged attention）、推测解码（speculative decoding）、分块预填充（chunked prefill）和前缀缓存（prefix caching）。在vLLM V0版本中，这些优化措施并非总能相互兼容。而要使它们无缝协同工作，需要对代码库进行大量修改，这些改动甚至深入到了内核层面。

为了进一步简化调度器相关的代码，V1版本还对请求批次的形成方式进行了显著调整。在V0中，vLLM的调度器总是形成“预填充批次”（仅包含需要预填充的等待请求）或“解码批次”（仅包含正在解码的运行请求）。但在V1中，调度器能够形成混合批次，其中可以包含处于任何状态的请求，如新的预填充、分块预填充、解码，甚至推测解码。

起初，唯一支持这些“V1批次”的注意力后端是FlashAttention软件包的CUDA版本。这意味着V1最初只能在NVIDIA的CUDA平台上使用，而AMD等其他厂商的GPU则无法获得支持。vLLM V0版本中用于AMD GPU的注意力内核是使用定制的C++/HIP代码实现的，这使得它们难以适应V1版本的新使用场景。有鉴于此，来自美国AMD、美国IBM研究院和美国红帽公司的团队决定共同开发一个基于Triton内核的vLLM V1新注意力后端，从而以平台可移植且更具开发者友好的方式，为AMD GPU提供V1版本的支持。

核心概念：理解注意力机制的关键术语

在深入探讨内核的工作原理之前，新媒网跨境获悉，理解一些基本术语至关重要。在vLLM中提及一个序列时，有三个关键量需要明确：

首先是上下文长度（context length），它表示已经计算过注意力并且键值（Key-Value，简称KV）张量已存储在分页KV缓存中的Token数量。

其次是查询长度（query length），它描述了本次调度迭代中需要计算的“新”Token数量。

最后是序列长度（sequence length），它是上下文长度与查询长度之和。

为了更好地理解上述概念，我们可以设想几种特殊情况：

1. 上下文长度为0： 这对应于“纯”预填充操作。在这种情况下，KV缓存中不包含与该序列相关的任何内容，必须从头开始计算提示词（prompt）的注意力。
2. 上下文长度大于0，查询长度约为1000个Token： 这是一个典型的“分块预填充”操作场景，我们将一个长提示词分割成较小的块进行处理。在这种情况下，新Token需要关注已存储在分页KV缓存中的旧Token。
3. 上下文长度大于0，查询长度为1个Token： 这是经典的“解码”操作。
4. 上下文长度大于0，查询长度约为3个Token： 这是进行推测解码（speculative decoding）时常见的模式。推测解码通常使用一个小型草稿模型来估计接下来的几个Token，然后由主模型进行验证。

一个适用于vLLM V1的Triton内核需要支持以上所有四种情况，以及介于它们之间的各种场景。

最初，唯一支持所有这些场景的Triton内核是prefix_prefill内核。该内核源自LightLLM项目，为vLLM V1在AMD GPU上运行提供了良好的起点。与本文中描述的所有注意力内核一样，它实现了平铺softmax算法（具体是FlashAttention-2论文中描述的算法）。在启动内核时，使用的网格大小为(batch_size, num_query_heads, query_length // BLOCK_M)，其中BLOCK_M是一个常数，设置为128。因此，prefix_prefill在批次中的序列之间、查询头之间以及查询维度中的“M块”之间并行处理任务。

每个程序负责沿着序列长度维度计算平铺softmax。本质上，它以“瓦片”（tiles）的形式遍历序列长度维度，其中序列维度中瓦片的大小由BLOCK_N（设置为64的常数）定义，查询长度维度中瓦片的大小由BLOCK_M定义。对于与上下文对应的瓦片，内核从分页KV缓存中读取KV张量。对于与查询对应的瓦片，KV张量从连续内存中读取（因为它们是在相同的正向传递中计算的）。图3展示了给定程序执行的工作。

尽管此内核支持上述所有四种情况，但美国IBM研究院、美国红帽公司和美国AMD的团队在运行标准性能基准测试后发现，V1的性能比V0基准性能慢约6倍。

Triton内核优化历程：三阶段突破

针对vLLM V1在AMD GPU上性能表现不佳的问题，美国AMD、美国IBM研究院和美国红帽公司的团队随即启动了一系列优化工作，力求提升其运行效率。这些优化工作按时间顺序分为三个阶段：

1. 前缀预填充内核的优化。
2. 针对解码序列采用的“拆分”方法，使用专用内核。
3. 一个统一的Triton内核，同时支持预填充和解码（以及介于两者之间的所有情况）。

接下来，我们将逐一探讨这些阶段的具体内容。

阶段一：前缀预填充内核的精细化调优

优化工作的第一阶段由美国AMD团队负责，旨在改善上述基准prefix_prefill内核的性能。

1. 线程束数量从8个减少到4个：
   prefix_prefill内核中存在较高的溢出率，严重影响了执行速度。美国AMD MI3xx系列GPU仅有512个向量通用寄存器（VGPRs），它们平均分配给特定SIMD上调度的线程束。如果内核需要更多寄存器（例如，处理过大的内存块），一些寄存器中的数据将被刷新到高带宽内存（HBM）中作为暂存空间，并在需要时重新用于即时执行，稍后溢出的数据再被恢复。8个线程束意味着2个线程束被调度到同一个SIMD上，因此每个线程束只有256个VGPRs。通过减少活动线程束的数量，团队消除了寄存器溢出。仅此一项改动就带来了约3-5倍的加速，并且是显著提升内核性能的最“唾手可得”的优化之一。

2. 参数化调度和内核配置以实现自动调优：
   第一步的优化减少了SIMD上线程束执行的并行性，也减少了线程束调度器在单个SIMD内交错获取、解码和执行指令的可能性。为了增加灵活性，需要重构内核以正确扩展单个程序处理的数据。团队开放了关键的内核启动参数（例如块大小、线程束数量、流水线阶段、循环展开、预取等），以便自动调优器可以探索针对序列长度、KV缓存页大小和硬件特性进行优化的配置。通过自动调优，团队为许多真实世界场景选择了最合适的配置。自动调优本身会因确定要调用哪个内核以及它是否已在缓存中而增加延迟。为了避免这种情况，决定直接将自动调优的参数发送给内核。

3. 对齐缓存循环中的块内存访问：
   分页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编译器能够最大程度地了解要加载的数据。这迫使编译器生成向量化的全局加载（global loads），从而提高了内存带宽效率。此外，这种对齐解锁了编译器可以利用的循环展开和流水线机会，进一步减少了内存密集型区域的开销。

4. 重构在线Softmax逻辑：
   在美国AMD团队的观察中，在线softmax的计算方式不够理想。它为P尺度和Acc尺度增加了额外的计算。这些计算发生在热循环内部，导致额外的寄存器使用和计算时间（向量操作是同步的）。团队重写了prefix_prefill内核中的softmax部分，以减少中间寄存器的使用（例如，重构累加、重新排序操作）。由此降低的寄存器压力减少了溢出风险，并改善了Triton生成的代码中的指令级并行性。

5. 精简内循环执行：
   作为优化工作的一部分，团队从不操作边界瓦片的内循环中消除了冗余的边缘情况条件检查。通过根据瓦片位置专门化循环体，非边缘循环现在可以在没有先前用于防止越界访问的条件分支的情况下执行。这减少了控制流分歧，并使编译器能够生成更高效的直线代码，从而提高了指令吞吐量和整体内核性能——特别是对于大型、对齐良好的工作负载。

阶段二：针对解码序列的“拆分”策略

优化工作的第二阶段始于美国红帽公司和美国IBM研究院团队注意到，prefix_prefill内核用于解码序列（上述第三种情况）的效率非常低。对于这些序列，查询长度始终为1个Token。如果我们思考瓦片在查询长度维度上的布局，可以立即看出，由于BLOCK_M设置为128，我们需要屏蔽掉瓦片中除了一个行之外的所有行，这意味着我们实际上进行了约100倍于实际所需的工作量。

虽然降低解码序列的BLOCK_M值似乎是显而易见的解决方案，但这里并非没有代价。因为vLLM在同一个批次中打包了解码和预填充请求，而BLOCK_M的值是一个编译时常量，因此批次中的所有序列必须使用相同的BLOCK_M值。这意味着，为解码序列优化BLOCK_M将对同一批次中的其他序列产生不利影响。

基于上述观察，美国IBM研究院团队首先尝试修改prefix_prefill内核，使其在查询长度为1时立即终止。这样，prefix_prefill仅用于计算批次中查询长度大于1的序列。随后，团队启动了第二个新编写的Triton内核——paged_attention_2d。该内核在查询长度大于1时立即终止，而对于查询长度为1的序列（即解码），它以高度优化的方式计算分页注意力。名称中的“2D”指的是用于启动内核的网格：它是一个形状为(batch_size, num_query_heads)的网格。值得注意的是，prefix_prefill使用了一个三维网格，其额外的维度对应于查询长度维度。由于解码操作不暴露这种额外的并行级别，因此不再需要它。

图5展示了此内核执行的工作。每个程序都在单个查询头上操作（在y轴上显示），并沿着上下文长度维度（在x轴上显示）执行平铺softmax算法。由于此内核仅处理分页KV缓存，因此团队将BLOCK_N设置为16个Token，以与vLLM中使用的默认页大小对齐。由于查询维度中没有冗余的“阻塞”操作，因此没有进行数量级上多余的工作。

新媒网跨境了解到，将此内核与prefix_prefill结合使用，使得吞吐量提高了约3.7倍。然而，对于使用分组查询注意力（Grouped Query Attention，简称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的倍数，可以确保矩阵乘法（matmuls）映射到矩阵核心，从而带来额外的性能提升。总的来说，这些与GQA相关的优化使吞吐量提高了约25%。

阶段三：统一注意力内核的演进

虽然此时的性能已得到显著改善，但“拆分”方法仍存在一些不足。具体来说，将工作拆分到两个内核（一个用于解码序列，另一个用于其他所有情况）有几个缺点。首先，需要维护两个复杂内核的代码。其次，在低延迟场景中，可能会受到基于CPU的启动开销的负面影响。最后，“拆分”解决方案没有考虑查询长度不完全为1，但可能仍然很小的推测解码情况。在这种情况下，使用prefix_prefill将再次变得非常低效。

为了克服这些限制，美国IBM研究院团队开发了unified_attention_2d Triton内核（参见图7）。该内核再次重新定义了启动网格的形状为(total_num_q_blocks, num_kv_heads)，其中total_num_q_blocks是一个与整个批次中查询长度总和成比例的量。每个程序沿着序列长度维度（图7中的z轴）执行平铺softmax算法。在查询头维度上进行阻塞，BLOCK_Q的大小恰好为QpkV（图7中的y轴所示），同时沿着“扁平化”批次维度（图7中的x轴所示）使用BLOCK_M = 16/QpkV进行阻塞。通过这种方法，团队在相同的矩阵乘法中捆绑了查询头和查询维度中的Token，这意味着在QpKV<16的情况下，可以向相同的张量核心指令中打包更多的工作。需要注意的是，扁平化批次维度中的阻塞（“Q块”）必须仔细定义，以确保每个Q块仅包含批次中一个序列的Token（请注意图7中每个查询末尾的一些Q块是如何被屏蔽的）。这个约束是必要的；否则，瓦片将需要同时处理多个序列，从而极大地复杂化逻辑。

现在，团队有了一个可以用于V1批次的单个内核，代码行数仅为239行，这使得长期维护变得容易得多。然而，美国AMD团队的实验表明，在某些场景下，Triton内核的性能仍然不足。特别是，美国AMD团队发现，为vLLM V0编写的C++/HIP解码内核在输出Token数量非常长时，仍然能提供更好的性能。团队确认，这是因为统一的Triton内核没有利用上下文长度维度上的任何并行性。

为了解决这个问题，美国IBM研究院团队在上述内核的基础上开发了一个扩展——unified_attention_3d，它遵循Flash-Decoding的思想，沿着上下文长度维度拆分工作，创建了额外的并行性，但代价是需要一个最终的归约步骤。3D内核在头部数量维度（num_kv_heads）或扁平化批次维度（tot_num_q_blocks）中没有足够的并行性来完全占用GPU时最为需要。这种情况通常发生在批次相对较小或批次中查询长度非常短（例如，纯解码）时。然而，如果批次非常大或批次包含几个非常长的查询，那么2D内核就具有足够的并行性，执行最终归约步骤的开销可能不值得。因此，美国IBM研究院团队实现了一个简单的启发式算法来决定何时使用2D内核与3D内核。

随后，美国AMD团队在统一的2D/3D注意力内核基础上贡献了额外的优化，使vLLM V1的性能达到了当前的水平。

1. SWA循环边界调整和配置更新： 优化了滑动窗口注意力（Sliding Window Attention，SWA）机制，将计算严格限制在活动窗口范围内。此前，内核会在整个序列长度上评估注意力分数，随后丢弃目标窗口之外的值，导致不必要的内存读取和计算开销。新实现引入了感知范围的索引和掩码逻辑，动态地将查询-键交互限制在相关的Token子集内。这项改动不仅降低了长序列场景下的二次计算成本，还改善了内存局部性、缓存效率和整体吞吐量——尤其对流式或增量解码工作负载有利。
2. KV和Q块的“cg”缓存修饰符（条件性地复用）： 应用缓存修饰符来优化全局内存加载——这在内存密集型负载的美国AMD GPU上尤其有利。
3. 2D注意力网格重排序以改进XCD映射： 重新设计了2D注意力网格布局，以增强空间局部性和跨计算芯片（Cross-Compute Die，XCD）的数据重用。新的映射策略重新排序了块索引，使得处理相邻查询-键瓦片的工作组优先被共同定位在同一个XCD上。这种对齐确保了共享重叠键/值区域的瓦片在相同的内存域上执行，显著提高了L2缓存驻留率并减少了跨芯片流量。通过最小化远程内存访问和改善共享数据的临时重用，这项优化提高了有效内存带宽，并带来了多芯片配置下更一致的延迟扩展。
4. 增加BLOCK_M以减少数据重载： 调整了瓦片高度（BLOCK_M），以更好地平衡寄存器利用率和内存访问效率。通过增加每个块处理的行数，每个线程组现在在驱逐之前，可以对加载的查询和键数据在更大范围的计算中进行重用，显著减少了冗余的全局内存读取。这项改动提高了计算与内存的比率，增加了算术强度，并提升了整体吞吐量——尤其对于内存带宽通常成为主要瓶颈的较长序列长度。新的配置还允许更好地重叠内存预取和计算，进一步提高了硬件利用率。
5. 专门化的预填充配置调优： 在拆分内核执行模型中，引入了专门针对预填充阶段工作负载的参数预设。这些预设自动调优了关键内核参数——如瓦片维度、线程束数量和流水线深度——以最大化GPU占用率和特定于预填充特性的数据重用。通过将瓦片大小和内存预取模式与预填充操作通常较大的批次和序列维度对齐，新配置提高了缓存利用率，减少了全局内存流量，并保持了高算术强度。
6. 修订内核选择逻辑（2D与3D）： 增强了vLLM的内核选择启发式算法，使其能更准确地根据当前序列长度、批次形状和Token分布，在2D和3D启动配置之间进行选择。更新后的逻辑考虑了工作负载几何结构和GPU占用率特性，确保每种启动策略都在其表现最佳的区域应用。这种自适应选择提高了各种输入模式下的内核效率——最小化空闲线程，改善内存访问规律性，并在小批次和大上下文工作负载中提供更一致的性能。
7. 扩展3D解码的网格拆分： 增加了沿第三个启动维度的分区程度，以生成更大的整体网格，从而提高GPU占用率和并行性。此调整允许更多的线程块并发调度，确保更好地利用可用的SM，特别是在单块计算可能限制并发性的长序列解码工作负载中。这项改动增强了设备间的负载平衡，并导致性能随序列长度扩展时更加一致。

性能基准测试：V1全面超越V0

本文详细阐述了数月以来为开源社区贡献的大量优化。每项优化的性能提升可以在文中链接的拉取请求中找到。在本节中，旨在评估这项工作的累积效应，并展示vLLM V1（基于Triton）在AMD GPU上真实的异构服务基准测试中，现在已超越vLLM V0（基于C++/HIP）。

团队将比较vLLM V0和V1在单个美国AMD MI300x GPU上运行mistralai/Mistral-Small-24B-Instruct-2501模型时的性能。所有实验均使用公开可用的Docker镜像和vLLM内置的基准测试工具进行，因此应易于复现。V0和V1分别使用了以下Docker镜像：rocm/vllm:rocm6.4.1_vllm_0.9.1_20250715和rocm/vllm:rocm7.0.0_vllm_0.10.2_20251006。

比较V0和V1的性能较为复杂，因为两个版本具有不同的默认设置。具体来说，控制vLLM可以并发处理的最大序列数的参数max-num-seqs，在V0中默认为256，而在V1中设置为1024。类似地，在V0中默认禁用分块预填充，而V1默认启用分块预填充，并且对于此特定模型，分块预填充参数max-num-batched-tokens默认为8192。这些参数会显著影响性能，并可能掩盖V0和V1之间与不同内核实现相关的差异。因此，团队运行了4种不同的实验来尝试控制这些影响，并隔离本文所述内核工作带来的益处。

1. V0（默认设置，除了禁用前缀缓存）：

   VLLM_USE_V1=0 vllm serve mistralai/Mistral-Small-24B-Instruct-2501 \
       --no-enable-prefix-caching
   

2. V0（增加最大序列数）：

   VLLM_USE_V1=0 vllm serve mistralai/Mistral-Small-24B-Instruct-2501 \
      --no-enable-prefix-caching \
       --max-num-seqs 1024
   

3. 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
   

4. 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吞吐量、首个Token生成时间（Time To First Token，TTFT）和Token间延迟（Inter-Token Latency，ITL）方面的性能。结果表明，即使在调整V0的默认参数以匹配V1之后，使用为vLLM V1开发和优化的Triton内核，在所有指标上仍能看到约10%的改进。

结语：展望未来与社区协作

本文详细介绍了美国AMD、美国IBM研究院和美国红帽公司团队如何通过编写优化的Triton内核并将其贡献给开源社区，为vLLM V1构建了优化的注意力后端。测试结果清晰地表明，基于Triton的vLLM V1在AMD MI300x GPU上的吞吐量比使用定制C++/HIP实现的vLLM V0高出10%。新媒网跨境认为，本文提供的细节和见解，对于正在为注意力或其他操作编写内核的Triton开发者而言，将是宝贵的参考。如果业界同仁对如何进一步提升vLLM性能有任何想法，欢迎随时沟通交流。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/vllm-v1-triton-kernel-boosts-amd-gpu-10-perf.html