vLLM 中的 DeepSeek-V3.2-Exp：细粒度稀疏注意力机制的应用

简介

我们非常高兴地宣布 vLLM 对 DeepSeek-V3.2-Exp 提供 Day 0 支持。该模型采用了专为长上下文任务设计的 DeepSeek 稀疏注意力（DSA）（论文）。在这篇文章中，我们将展示如何在 vLLM 中使用该模型，并深入探讨在 vLLM 中支持 DSA 时遇到的挑战。

特别是，DSA 的闪电索引器（Lightning Indexer）结合稀疏注意力机制，给持续批处理（Continuous Batching）和分页注意力（Paged Attention）带来了挑战。例如，我们需要为索引器模块分别处理预填充（Prefill）和解码（Decode）阶段，并仔细管理不同的缓存布局。

在性能方面，vLLM 集成了 DeepGEMM 中的闪电索引器 CUDA 内核，以及 FlashMLA 中的新稀疏注意力内核。我们还非常兴奋地宣布 Blackwell 支持。通过与 NVIDIA 的合作，您可以直接在 B200 和 GB200 上运行该模型！

使用指南

要开始使用 DeepSeek 3.2，请按照使用指南（recipes）中的安装说明进行操作。我们正在通过此 PR 继续改进初步支持。有关已知问题，请参阅跟踪议题。

安装完成后，在 16×H100、8×H200 或 8×B200 上，您可以使用张量并行运行该模型（专家并行目前存在一个小问题，我们正在修复）。

vllm serve deepseek-ai/DeepSeek-V3.2-Exp --tensor-parallel-size 8

为了进行大规模部署，我们期待在本周晚些时候分享使用 llm-d 实现的一键式 Kubernetes 部署方案。该方案利用 NIXL 进行 vLLM 的 PD 解耦，然后针对每个 P 和 D 实例，高效地将请求路由到不同的数据并行秩（data parallel ranks）。文档即将发布。

引擎启动后，建议您使用长输入或期望长输出的提示词进行测试。我们推荐将其与 V3.1-Terminus 进行比较，因为二者在相同的数据混合上进行了持续预训练。

我们目前仍在核实 vLLM 的实现与官方精度结果的一致性。在之前的模型权重版本中，我们的 GSM8K 和 GPQA-Diamond 分数达到了预期水平，并显示其表现与 V3.1-Terminus 相似。

在 vLLM 中实现 Top-K 稀疏注意力

新的缓存条目与量化方案

闪电索引器模块具有专门用于索引的缓存 K 值。这意味着对于每个 token，现在都有另一个用于索引器的 K 缓存。vLLM 分配了独立的缓冲区来保存索引器 K 缓存，与 MLA K 缓存隔离开来。

另一个值得关注的重点是 FP8 KV 缓存的处理，该模型对此提供了支持。对于 MLA，每个 token 的 KV 缓存为 656 字节，结构如下：

• 前 512 字节：“量化 NoPE”部分，包含 512 个 float8_e4m3 数值。
• 随后的 16 字节：缩放因子，包含 4 个 float32 数值。第一个 float32 是前 128 个 float8_e4m3 数值的缩放比例，第二个是接下来 128 个的缩放比例，依此类推。
• 最后 128 字节：“RoPE”部分，包含 64 个 bfloat16 数值。为了保证精度，这部分未经过量化。

然而，对于索引器 Key 缓存，它是按块（block）存储的。这也是我们为何仅支持该模型块大小为 64 的原因之一；另一个原因是 FlashMLA 也适配了此尺寸。前 block_size * head_dim 个条目包含数值，其余部分包含缩放因子。

x_fp8[ :, : block_size * head_dim] = x_scaled.view(num_blocks, block_size * head_dim).view(dtype=torch.uint8) 
x_fp8[ :, block_size * head_dim :] = scales.view(num_blocks, block_size).view(dtype=torch.uint8)

在索引器中，单个 token 的缓存并不是连续存储的。

带有掩码的新计算方式

对于每个新的查询 token，它现在通过索引器计算出 Top-2048 个需要关注的 token。token 的查询是一个形状为 (h, d) 的张量，其中 h 是查询头数，d 是头维度。长度为 n 的上下文是一个形状为 (n, d) 的 2D 张量。计算出的 Logits（查询与上下文之间的相关性得分）是一个形状为 (n, h) 的张量。通过形状为 (h,) 的头权重对 Logits 进行加权，我们得到一个形状为 (n,) 的张量。我们需要生成一个包含 Top-2048 个 token 索引的 (2048,) 整数张量，如果不足 2048 个 token，其余部分填充 -1。

虽然理解单个查询 token 如何选择关注的索引很简单，但批处理的情况更为复杂。我们来拆解一下。

新的 DeepGemm 函数调用方式如下：

logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)

对于来自同一请求的多个查询 token（长度为 q，即预填充情况），它们存储在形状为 (q, h, d) 的张量中。上下文仍有 n 个 token，因此上下文仍然是一个形状为 (n, d) 的 2D 张量。Logits 是一个形状为 (q, n, h) 的张量。通过头权重对 Logits 加权后，得到形状为 (q, n) 的张量。我们需要生成一个形状为 (q, 2048) 的整数张量，表示 Top-2048 个 token 的索引。由于因果关系，每个查询 token 只关注其之前的 token。我们需要标记每个查询 token 的上下文起始和结束位置。我们使用 ks 标记上下文开始，ke 标记上下文结束。ks 和 ke 都是 (q,) 形状的整数张量。在这种情况下，ks 全为 0，而 ke 将是 list(range(n - q, n, 1))。

最后，考虑如何批量处理多个请求。我们有 b 个请求，每个请求有 q1, q2, ..., qb 个查询 token，以及 n1, n2, ..., nb 个上下文 token。查询 token 将被批处理成形状为 (q1 + q2 + ... + qb, h, d) 的张量。上下文将被批处理成形状为 (n1 + n2 + ... + nb, d) 的张量。Logits 将被批处理成形状为 (q1 + q2 + ... + qb, n1 + n2 + ... + nb, h) 的张量。我们需要生成一个包含 Top-2048 个索引的 (q1 + q2 + ... + qb, 2048) 整数张量。

我们需要标记每个查询 token 的起始上下文和结束上下文。我们使用 ks 标记起始上下文，ke 标记结束上下文。ks 和 ke 都是形状为 (q1 + q2 + ... + qb,) 的整数张量。

在这种情况下，ks 将是 [0] * q1 + [q1] * q2 + ... + [q1 + q2 + ... + qb] * qb。这里 * 表示重复列表。ke 将是 list(range(n1 - q1, n1, 1)) + list(range(n2 - q2, n2, 1)) + ... + list(range(nb - qb, nb, 1)) 加上 ks 的偏移量。

计算出 Logits 后，我们需要执行 topk 操作。然而，一个明显的挑战是，在高批处理量和长上下文情况下，Logits 张量在运行逐行 topk 之前就会被实例化。

融合算子、更多内核及 Blackwell 支持

当我们开始优化性能时，首先处理了一些低垂的果实：

• Top-K 可以通过融合内核来表达，以获得更好的性能。DeepSeek 团队的 TileLang 内核是一个很好的参考！
• 我们在 MLA 隐向量和索引器 Key 向量写入 vLLM 页表时使用了量化。事实证明这并不简单，因为正如我们之前解释的，该量化方案是全新的且与众不同。

我们也非常高兴地宣布为该模型提供开箱即用的 Blackwell 支持。我们致力于使 Blackwell 平台在未来的模型发布中成为一等公民，因为它的效率有助于实现最佳性能！

进行中的工作

我们目前对 vLLM 中 DSA 及相关稀疏注意力的优化还仅仅触及了表面。在接下来的几周内：

• 我们计划将支持的架构扩展到 Hopper 和 Blackwell 之外。
• 我们将扩展支持到 AMD 和 TPU 等其他硬件。借助 vLLM 的可扩展系统，开发人员可以直接添加对模型的支持。例如，vllm-ascend 和 vllm-mlu 已经支持了 DeepSeek V3.2！
• 我们正在持续测试大规模宽专家并行（Wide EP）服务和解耦。
• 您很快就能使用该模型运行端到端的 RL 循环。
• 我们将探索 DeepSeek 提出的“用于短序列预填充的掩码 MHA 模式”。
• 在此版本中，我们移除了 Hadamard 变换，因为观察到它对精度没有影响。我们将进一步研究！

致谢

vLLM 社区中的以下团队参与了该模型的支持工作：

• vLLM：Chen Zhang, Yongye Zhu, Kaichao You, Simon Mo, Zhuohan Li
• Red Hat：Lucas Wilkinson, Matt Bonanni, Wentao Ye, Nicolo Lucchesi, Michael Goin, Robert Shaw, Tyler Michael Smith
• Meta：Lucia Fang, Xiaozhu Meng, Lu Fang
• NVIDIA：Ray Wang, Barry Kang, Daniel Campora, Julien Demouth, Siyuan Fu, Zeyu Wang, Pen Chun Li

作为 vLLM 团队，我们要感谢 DeepSeek 团队开源该模型、技术和内核，并感谢 DeepSeek 领导层对 vLLM 的信任与支持！