跳到正文
vLLM 官方博客· Inferact and the vLLM Team·· 21 小时前精选AI 评分63

vLLM 详解 DeepSeek-V4.1-Flash 优化:Agent 场景吞吐提升 5 倍

DeepSeek-V4.1-Flash on vLLM: 5x Agentic Throughput Since Day 0

AI 导读

Inferact 与 vLLM 社区在 DeepSeek-V4.1-Flash 发布三周内完成优化,低并发速度提升 1.9 倍,150 TPS 约束下吞吐提升 5.3 倍。

推荐理由

vLLM 团队拆解了 SWA bounded replay 和一系列内核优化,说明五倍吞吐提升具体来自哪里,方法可复用于其他 Agent 服务场景。

正文 · AI 翻译

TL;DR:在 DeepSeek-V4.1-Flash 发布后的三周内,Inferact 和 vLLM 社区对该模型进行了优化,在低并发下实现了 1.9 倍加速,在 150 TPS 约束下吞吐量提升了 5.3 倍。性能提升来自:

  • 我们实现了带 CUDA graphs 的 SWA 有界重放,将 TTFT 降低了约 30%。

  • 我们集成了 DeepSeek 新发布的内核,包括 MegaAttention、Mega-mHC、Mega-Gate 和 DeepSelect。

  • 我们激进地融合并并行化了剩余的内核,包括 mHC 侧流,并将 all-reduce 与其前后操作融合为单个内核。

DeepSeek V4.1 为长时程智能体服务任务引入了一种高效架构:凭借其因果编码器-解码器(CED)架构,模型在解码时每个 token 激活 16B 参数,但在预填充时仅激活 8B 参数。该模型还极其节省内存。它结合了多种技术来缩小 KV 缓存大小:压缩稀疏注意力 2(CSA2)、FP4 KV 缓存以及层间 KV 缓存共享,将全局 KV 占用降至每 token 890 字节。本文展示了我们如何将 DeepSeek 的这些模型级优化与 vLLM 侧的系统优化相结合,在 SemiAnalysis AgentX 智能体服务基准上实现 5 倍吞吐量。我们重点介绍两类优化:SWA 有界重放和内核相关优化。

SWA 有界重放

DeepSeek-V4.1-Flash 保留两种 KV 缓存。全局 KV 经过压缩、跨层共享并以 FP4 存储,每 token 约 890 字节(V4.1 报告)。滑动窗口(SWA)KV 未压缩,以 FP8 存储,覆盖 40 层中每层的最后 128 个位置。

SWA KV 带来两项成本:

  1. 前缀缓存必须在每个可能的命中边界存储它,存储成本超过全局 KV 的 10 倍。

  2. 预填充在每个提示 token 上运行第 21–39 层,而解码仅读取它们的最后 128 个位置。

一种直接的方法是重新计算 SWA KV 而不是缓存它。然而,精确重计算代价高昂,因为每层的 128 token 窗口依赖于下一层中更早的位置,因此跨 L 层重建它意味着重放大约 L × 128 个 token。

DeepSeek V4.1 引入了 SWA 有界重放,以精确性换取效率。它仅重新运行最后 128 个 token,并在重放起点处裁剪 SWA 窗口。结果并非逐位精确,但 DeepSeek 报告质量损失可忽略不计(详见下文)。vLLM 在两处应用它,分别对应一项成本。

编码器侧:在缓存命中时重建窗口

通过编码器侧重放,vLLM 仅缓存全局 KV 并跳过 SWA KV。在长度为 H 的前缀命中时,它重新运行 token [H − 128, H) 以重建 SWA KV,窗口在 s = H − 128 处裁剪。

解码器侧:跳过大部分提示预填充

在 DeepSeek V4.1 的 CED 架构中,第 20 层计算解码器的全局 KV,第 21–39 层复用它。因此,vLLM 在每个 token 上运行第 20 层以生成该全局 KV,而仅在每个请求的最后 128 个 token 上运行第 21–39 层。对于长提示,这跳过了近一半的模型。

为裁剪层使用 CUDA graphs

裁剪后,第 21–39 层在 GPU 上的计算量极小,若以 eager 模式运行,内核启动开销将占主导,GPU 会处于空闲状态。它们的输入形状也与第 0–20 层不同,因此这两部分无法被捕获到同一个 CUDA graph 中。vLLM 的可中断 PIECEWISE graph 本就会在模型中间断开,这提供了一个天然的切分点:第 0–20 层在完整批次上捕获,第 21–39 层则在裁剪后的批次上单独捕获。这使得 CUDA graph 可用于裁剪后的 prefill,并让第 21–39 层使用各自的捕获尺寸,从而获得更好的 graph 覆盖。

SWA 有界重放(bounded replay)在 DeepSeek-V4.1 上默认开启,由 --[no-]swa-bounded-replay 控制。

精度与性能结果

尽管 SWA 有界重放并非精确计算,但 DeepSeek 报告称质量损失可忽略不计。我们在 vLLM 上通过包括 GSM8K 和 GPQA 在内的基准测试确认了这一点,未观察到有意义的精度差异(差距约在 1.5 个标准误以内)。

在性能方面,编码器侧每次命中会以缓存空间换取一个窗口的 prefill,因此加速主要来自解码器侧。我们测量了三种设置下的单请求 prefill TTFT:重放关闭;重放开启但不使用解码器 CUDA graph(第 21–39 层以 eager 模式运行,且仅裁剪 eager 步骤);以及重放开启并使用解码器 CUDA graph。

使用 CUDA graph 的解码器重放将 prefill 计算时间缩短了 30–40%。 CUDA graph 对短提示词最为关键,此时内核启动是瓶颈:若不使用它们,启动开销会超过 GPU 节省的时间,重放反而比基线更慢(在 DEP2 上 1K 时最多慢 12%)。对于长提示词,GPU 计算量足够大,可以掩盖启动开销,因此 eager 重放已能获得大部分收益,CUDA graph 再额外提升几个百分点。

内核

DeepSeek 在发布 DeepSeek-V4.1-Flash 的同时,在其三个代码仓库中发布了新内核。DeepSelect 是用于 DeepSeek Sparse Attention 的新 top-k 库。DeepGEMM 新增了稀疏索引器内核以及若干与 GEMM 相关的融合内核。FlashMLA 新增了 NVFP4 KV 缓存支持和融合注意力内核 MegaAttention。我们已将其中若干开源内核集成到 vLLM 中,并在 #57448 中跟踪进展。

Mega-mHC(#56962)。 Mega-mHC 将 mHC 链融合为一个内核:post 步骤、延迟 pre 步骤以及 RMSNorm。它取代了现有的 TileLang 融合路径,而 DeepGEMM 的实现现已优于后者。在 NVIDIA GB200 上,该内核比 TileLang 版本快 1.14–1.51 倍。

Mega-Gate(#56266)。 Mega-Gate 将 MoE 路由器(gate GEMM、专家打分、偏置和 top-k 选择)融合为一个内核。此前,这些操作以一个 GEMM 后接单独的 top-k 内核运行,额外产生一次启动开销以及分数在内存中的往返。这一融合在中等批次大小下带来 1.18–1.31 倍的内核加速。

mHC 多流重叠(#57603)。 在 V4.1 中,mHC 系数偏移了一个子层,因此下一个接缝的系数 GEMM 只读取在注意力或 FFN 运行之前就已存在的残差流。在下一个 post/pre 步骤将它们合并之前,双方都不需要对方的输出。在小批次大小下,vLLM 现在会在侧 CUDA 流上计算下一个 mHC 块的系数,与注意力和 FFN 并行执行。这隐藏了原本会处于延迟受限解码关键路径上的工作。在 TP4 低延迟场景中,这将延迟降低了约 4%。

稀疏 MQA logits(#56254)。在 V4.1 中,后续的 indexer 层从固定的 16K 候选位置集合中选取其 top-k。此前的实现会计算整个上下文的分数,并在打分前屏蔽所有非候选块。DeepGEMM 的稀疏 kernel 仅对候选位置打分,因此开销不再随上下文增长。在 NVIDIA GB300 上,每层在 8K token 时快 1.2×,在 512K 时快 14–23×。在 4× NVIDIA GB300 上端到端,decode 提升 3–6%。Prefill 在 512K 时快 1.43×,在 1M 上下文时快 2×。

采用 NVFP4 压缩 KV 的 MegaAttention(#56935)。FlashMLA 的 MegaAttention kernel 在单次启动中完成 query RoPE、稀疏注意力、对输出做逆 RoPE 以及 FP8 转换,直接写入输出投影读取的缓冲区。这消除了注意力与下一层之间单独的 kernel 和内存往返。它还读取一种新的 NVFP4 压缩 KV 格式,比之前的 FP8 KV cache 小 45%。MegaAttention 还通过激进的融合消除了操作之间的 HBM 写入,将 kernel 效率提升 1.45×。

低延迟融合 WO-A kernel(#58634)。针对 Blackwell 上的小批量 decode,我们将逆 RoPE、FP8 量化、WO-A 批量 GEMM 和 MXFP8 重新量化融合进单个 CuTe-DSL kernel,将 WO-B 之前的链路从三个 kernel 减少为一个。关键思路是将中间激活保留在片上,并将数据搬运与计算流水线化,避免在小批量下占主导的额外 kernel 启动和全局内存往返。这将融合 WO-A 路径提升最多约 2.1×,并在低并发下将 inter-token 延迟降低最多约 6–7%。

Engram。V4.1 的 Engram 层根据哈希后的 token n-gram 从两个大型 FP8 表中查找行。每一步只读取少数几行,因此表的放置和查找延迟比计算更重要。我们异步预取 CPU 卸载的 Engram 查找,将主机内存访问与 decoder 计算重叠,以加速低批量 decode(#56512)。Engram head 采用统一的 TP/DP 方案进行分片,同址的 DP 副本共享相同的主机表,避免冗余副本以及查找路径上的任何 DP 通信(#57651)。对于这些大型主机驻留表,我们还支持透明大页(THP)以减少缺页开销,为 prefill 带来最高 10× 的查找 kernel 加速(#56926)。我们还针对可用大页不足的情况添加了优化(#59327)。

Agentic 性能

我们使用 SemiAnalysis AgentX 基准作为代表性的 agentic 服务负载来衡量性能(详见我们之前的文章)。这些优化共同使 vLLM 相比我们 day-0 的实现取得了显著的性能提升。如图 6 所示,我们的低延迟结果相比 day-0 结果提升 1.9×,高吞吐结果提升约 5×。

对于低延迟服务,我们使用 TP4 搭配 FlashInfer 注意力。小批量 decode 主要受内存带宽限制,因此将模型权重分片到四块 GPU 上很合适。我们也尝试了 MegaAttention,但其主要优势在于融合有更大发挥空间的高吞吐场景。在 TP4 下,这一收益要小得多,在我们的运行中 FlashInfer 最终更快。

为了获得高吞吐量,我们切换到 DEP2,使用 DP attention 并将专家拆分到各 GPU 上。由于 V4.1 在所有 head 之间使用共享的 KV latent,TP 会在各 GPU 上重复 KV cache。DP 避免了这种重复:每个 GPU 只为其所服务的请求存储 KV,并通过会话亲和性在多轮对话中保持 prefix-cache 的局部性。MegaAttention 进一步通过 NVFP4 削减了每个请求的 KV 占用,相比 FP8 几乎减半,并提升了每 GPU 的并发数。

值得注意的是,V4.1 的内存效率非常高,在整个基准测试过程中都不需要 KV cache offloading。我们预计在更高并发度下,配合 P/D 分离,KV cache offloading 将开始发挥作用。

SWA 有界重放,加上 prefill 侧的内核优化,也大幅改善了 TTFT。

图 7 展示了优化后的 TTFT–吞吐量权衡。在约 100K 吞吐量下,通过三项优化相结合,TTFT 下降了近 70%:

  • SWA 有界重放让模型的上半部分只处理最后 128 个 token,将 prefill 计算量大致减半。

  • CUDA graphs 让经过裁剪的小型重放在 GPU 上保持高速运行,而不受 CPU 内核启动的制约,从而让我们充分实现加速。

  • 内核改进加速了模型计算。

致谢

我们感谢 DeepSeek 开源 DeepSeek-V4.1-Flash 及相关内核,感谢 Inferact 团队最初的模型启动与优化工作,感谢 NVIDIA 的合作与支持,以及感谢 SemiAnalysis 提供 AgentX 基准测试。

来源:vLLM 官方博客 · vllm.ai