简介与议程
大型语言模型目前正面临注意力机制（Attention）的扩展极限。即使采用高度优化的实现，KV 缓存内存也会随着序列长度线性增长，而预填充（prefill）延迟则呈二次方增长。随着工作负载对 128k token 及以上上下文的需求日益增长，我们需要新的架构方法。
诸如 Qwen3-Next、Nemotron Nano 2、MiniMax-Text-01 和 Granite 4.0 等混合模型将注意力机制与 Mamba 或线性注意力（linear attention）等替代方案相结合，为未来指明了方向。通过混合多种机制，它们既保留了建模质量，又实现了高效的长序列推理。vLLM 社区近期的工作已将混合模型从 V0 版本中的实验性方案提升为 V1 中全面支持的一等公民。
本文将解释这一转变是如何发生的、其重要性所在，以及开发者在 vLLM V1 上部署混合模型时可以期待的性能提升。
[图 1：从 vLLM V0 迁移到 V1 后，代表性混合模型 granite-4.0-h-tiny（7B 参数，1B 激活参数）的吞吐量提升对比]
注意力机制就是一切？
注意力机制已成为大语言模型的基石，在捕捉自然语言依赖关系方面依然非常有效。vLLM 及其依赖项（包括 FlashInfer 和 FlashAttention）等推理引擎已在现代 GPU 上对注意力机制进行了深度优化。这些优化包括分页 KV 缓存（paged KV cache）、分块（tiled）及在线 Softmax 内核、Tensor Core 的深度利用、张量内存加速器（TMA）的集成，以及通过量化来压缩内存和提高吞吐量的广泛应用。
尽管取得了这些工程进展，但当扩展到极长序列时，注意力机制面临两个根本性限制。首先，KV 缓存随序列长度和批处理大小线性增长。每生成一个 token 都会附加新的 Key 和 Value 向量，当序列扩展到数十万个 token 时，该缓存很快会成为 GPU 上最大的内存占用者之一。其次，预填充阶段（以首个 token 生成时间 TTFT 衡量）随提示词长度呈二次方增长。对于 128k 或更长的提示词，预填充延迟会急剧增加，有时甚至导致推理变得不切实际。
注意力机制在许多方面依然是顶尖的，但这些扩展瓶颈是研究人员和从业者探索混合模型的主要原因。
为什么长序列如此重要
长序列并非抽象的研究目标，而是现实世界工作负载的要求。
检索增强生成（RAG）就是一个典型的例子。在这种模式下，用户的查询会与多个检索到的文档配对，并将这些文档附加到提示词中。根据检索到的文档数量及其大小，提示词长度很容易从几千个 token 激增到数万甚至数十万。
另一个驱动力来自智能体（agentic）模式。在此类工作负载中，模型在循环中运行，生成输出、与工具或系统交互，然后将结果整合回上下文中进行进一步推理。每一轮交互都会延长提示词，随着工具调用次数的增加，序列长度可以迅速达到非常大的规模。
最后，推理模式也明确鼓励使用长序列。指示模型“逐步思考”或生成中间推理链的提示词会将所有中间 token 保留在上下文中。这种结构化推理提高了准确性，但大大扩展了提示词的大小。
[图 2：RAG 查询、智能体工具循环和推理追踪的示例，展示了上下文长度如何以不同方式扩展。]
在这些场景中，混合模型旨在通过结合注意力层与在内存和时间上扩展性更好的替代方案（如 Mamba 或线性注意力）来提供高效性。
状态空间模型简史
状态空间模型（SSMs）历史悠久且底蕴深厚，在控制理论和动力系统等领域已取得成功。虽然它们历史上是在连续时间域中定义的，但它们也拥有一种离散化形式，可以应用于我们处理 LLM 推理时常见的 token 序列。
我们将从 2021 年提出的经典 SSM——S4 的离散形式开始。该算法如图 3 所示。S4 通过内部潜在状态 h 将长度为 T 的输入序列 x 映射到输出序列 y。与需要二次方计算量进行预填充的注意力机制不同，S4 在序列长度上是线性的，对于长度为 T 的序列仅需 T 次更新。同样重要的是，潜在状态 h 具有固定的维度 N，不会随序列长度增长。这使得 S4 在表示上非常高效，尤其是在 KV 缓存成本激增的长序列场景下。
[图 3：S4 循环图，展示了包含矩阵 A、B、C 和潜在状态 h 的每时间步更新。]
然而，S4 在需要选择性复制 token 或上下文推理的任务上表现不佳。虽然循环结构捕捉到了依赖关系，但它没有提供足够的灵活性来处理这些能力。
2023 年 Mamba-1 的出现通过允许矩阵 A、B 和 C 在每个时间步发生变化，扩展了 S4 的定义。这引入了选择性关注不同 token 的能力，极大提升了选择性复制和推理任务的性能。尽管如此，Mamba-1 的推理性能在现代 GPU 上仍落后于注意力机制，因为其并行实现并未充分利用矩阵乘法，无法有效地发挥 Tensor Core 的性能。
随着 Dao 和 Gu 在 2024 年发布里程碑式的 Mamba-2 论文，Tensor Core 的问题得以解决。该论文表明：(a) SSM 实际上可以被定义为从输入序列到输出序列的矩阵变换（而不是图 3 中的 for 循环）；(b) 在矩阵 A 中引入更多结构可以带来极高效的实现；(c) Mamba-2 实际上等同于另一种注意力变体：线性注意力。
线性注意力由 Katharopoulos 等人 在 2020 年提出。作者证明，将 Softmax 近似为内核特征映射的线性点积，会引出一种不同类型的注意力机制：不再随序列长度呈二次方扩展。这一发现引发了大量研究，并导致了该核心思想的许多变体。最近一些值得注意的线性注意力变体包括 Lightning Attention 和 Gated Delta Net，它们分别被 Minimax-Text-01 模型和 Qwen3-Next 模型所使用。
混合模型：vLLM V1 中的格局
随着这些机制的就位，领域内涌现了大量将完全注意力机制与 Mamba 或线性注意力变体相结合的混合架构。逻辑很明确：保留完全注意力机制以利用其建模优势，同时插入 Mamba 或线性层以提高规模化效率。
[图 4：vLLM V1 支持的混合模型表，突出显示了注意力机制、Mamba 变体、线性注意力和 MoE 的使用。]
在图 4 中，我们提供了目前 vLLM V1 支持的完整混合模型列表，并指出了每个模型使用的 SSM 或线性注意力机制类型。这些模型表明，混合化并非小众实验，而是各组织间积极的设计选择。
混合模型的状态管理
在 vLLM 中支持混合模型需要对状态进行精细化处理。注意力层依赖于分页 KV 缓存（每块 16 个 token，约占 64 KiB）。相比之下，Mamba 层为每个序列维护一个巨大的固定大小状态（约 2.57 MiB），该状态是在原位更新的，而不是通过追加的方式。
[图 5：注意力机制中的 KV 缓存块，显示 16-token 的组，每组约 64 KiB。数据基于 NVIDIA-Nemotron-Nano-12B-v2 计算]
[图 6：Mamba 状态结构，显示每个序列一个大的原位更新状态。数据基于 NVIDIA-Nemotron-Nano-12B-v2 计算]
尽管单个 KV 缓存块比 Mamba 状态小，但这种关系在长上下文中会发生逆转。对于 128k 的序列，KV 缓存的大小可能是 Mamba 状态的近 200 倍，这说明了混合方法为何如此具有吸引力。
vLLM V0 中的混合状态管理
在 vLLM V0 中，混合支持是通过一种实用但脆弱的 hack 实现的。KV 缓存分配作为块高效管理，但 Mamba 状态是单独分配的。每个 Mamba 层为每个活动序列维护一个状态张量，其大小由用户定义的参数 max_num_seqs 决定。
[图 7：V0 架构，显示用于注意力机制的分页 KV 缓存和单独的每个序列 Mamba 状态张量。]
这种方法迫使用户猜测该参数的正确值以避免崩溃。设置得过高有 CUDA 内存溢出的风险，而设置得过低则会降低并发性。这导致了糟糕的可用性，并引发了使用新混合模型的开发者的频繁抱怨。
vLLM V1 中的统一状态管理
在 V1 中，混合模型支持是围绕一个能够同时管理 KV 缓存和 Mamba 状态的统一分配器重建的。这种设计不仅优雅，还启用了前缀缓存（prefix caching）、KV 缓存传输和预填充/解码分离（prefill/decode disaggregation）等高级功能。通过将混合模型支持纳入 V1，它们现在也可以受益于诸如 torch.compile 和改进的调度等新优化。
但统一状态管理对于基于 Mamba 的混合模型并非开箱即用。虽然 V1 中已经通过混合内存分配器存在对“混合”状态的支持，但它仅支持像 Gemma 3、Llama 4 和 gpt-oss 这样将完整注意力层与滑动窗口注意力层（SWA）混合的模型。
[图 8： V1 中 KVCacheGroups 的示例，其中相似层的组共享缓存张量以简化内存管理。]
图 13 展示了对于具有完整注意力 (A) 和滑动窗口注意力 (SWA) 的模型，混合内存分配器的行为。分配器将相同类型的层分组为 KVCacheGroups。每个 KVCacheGroups 的状态以交叉方式存储在多个 KVCacheTensors 中，形成一个块。关键点在于，对应于不同组的块共享相同的 KVCacheTensors。为了便于简单的内存管理，我们必须确保块大小（以 token 为单位）以及页大小对于所有 KVCacheGroups 都是相同的。
适配 Mamba 分配器
Mamba 使情况变得复杂，因为其状态页的大小远大于注意力块。为了统一它们，我们需要放宽所有组必须使用相同块大小的要求。相反，注意力块的大小会自动增加，直到它们与 Mamba 的页大小对齐。同时，Mamba 页会被稍微填充，以使页大小完全相等。这种页对齐过程在图 9 和图 10 中针对代表性模型 NVIDIA-Nemotron-Nano-12B-v2 进行了说明。
[图 9：页对齐前 – 注意力页大小远小于 Mamba 页大小]
[图 10：页对齐后 – 注意力页大小和 Mamba 页大小完全相等]
虽然不同寻常的块大小（例如每个注意力块 672 个 token）看起来可能效率低下，但实证测试显示对性能的影响很小，这可能是因为当 Mamba 或线性层占主导地位时，注意力机制对运行时的贡献较小。
虽然一些注意力内核支持任意块大小（例如 Triton 注意力后端提供的内核），但另一些（如 FlashInfer 提供的 TRT-LLM 内核）则不支持。为了克服这一限制，我们最近实现了将 KV 缓存管理所使用的块大小与内核所见的块大小解耦的功能。这使得在以干净方式管理混合状态的同时，仍能使用 TRT-LLM 内核在 Blackwell GPU 上获得最佳性能。
追求极致
在调试初步的 V1 集成时，我们发现了一个微妙的问题，它源于 Mamba KVCacheGroups 与注意力 KVCacheGroups 共享相同的 KVCacheTensors。每个组对同一个张量使用不同的视图，且这些视图看起来截然不同，如图 11 所示。
[图 11：更改步幅（strides）前：注意力视图和 Mamba 视图不兼容]
每个注意力页包含 Key 张量和 Value 张量，在使用 FlashInfer 后端时，Key 和 Value 在 GPU 显存中以块对块的方式交叉存储。另一方面，Mamba 页包含卷积状态（Conv）和状态空间模型状态（SSM），默认情况下它们*不是*交叉的。因此，我们存储所有块的 Conv 状态，然后是所有块的 SSM 状态。这意味着通过 Mamba 视图写入一个块会导致注意力视图中另一个块的数据损坏。
为了解决这个问题，我们改变了 Mamba 状态张量的步幅方式，以确保两个视图完全兼容。这种方法如图 12 所示。我们可以看到注意力块和 Mamba 块现在完美对齐，不会发生数据损坏。
[图 12：更改步幅后：注意力视图和 Mamba 视图兼容]
某些注意力后端（如 FlashAttention）实际上并不以图中所述的块对块方式存储 Key 和 Value。然而，我们发现可以再次更改步幅，但这次是针对注意力视图，以确保 FlashAttention 也可以用作混合模型的后端。FlashAttention 是 vLLM 在许多情况下的默认注意力后端，因此这是淘汰 V0 的又一个重要步骤。
性能工程以淘汰 V0
为了完全淘汰 V0 的混合模型支持，还需要进行其他几项更改。大量的建模代码被重构，并进行了一系列性能优化，以确保我们能够超越 V0 的性能。
特别是，许多 Mamba 和线性注意力变体是使用 Triton 内核实现的。Triton 非常适合开发这些新模型，因为它提高了生产力并在不同平台上表现良好，但它在 CPU 启动开销方面有一个已知问题。这些开销在低延迟场景（例如小批次的 inter-token 延迟和/或激活参数较少的模型）中会严重影响性能。
为了克服这个问题，我们在 Mamba 后端中实现了分阶段的 CUDA Graphs 支持：从即时执行开始，到分段图，再到纯解码批次的完整图，最终实现了一个结合了解码器完整图和混合批次分段图的混合方案（FULL_AND_PIECEWISE）。这一最终优化使我们能够恢复 V0 的性能，并在某些场景下显著超越它。通过默认启用 FULL_AND_PIECEWISE，我们终于可以默认启用 V1 混合模型支持，并开始剥离 V0 代码。
最后，前缀缓存现在也支持使用 Mamba-2 的混合模型，尽管它目前处于实验状态。如果您有兴趣，请尝试使用，如果遇到问题，欢迎提出 issue。
基准测试设置
基准测试在 vLLM v0.10.2 上执行，这是最后一个同时支持 V0 和 V1 混合模型的版本，允许进行直接比较。服务器运行在 H100 GPU 上，客户端生成 32k 长度的随机输入和 128-token 的输出，并对并发级别进行扫测。
vllm serve $MODEL --trust-remote-code vllm bench serve \ --model $MODEL \ --dataset-name random \ --random-input-len 32768 \ --random-output-len 128 \ --max-concurrency $MAX_CONCURRENCY \ --num-prompts $((10 * $MAX_CONCURRENCY)) \ --ignore-eos --seed $MAX_CONCURRENCY
我们对两个模型进行了基准测试：NVIDIA Nemotron-Nano-12B-v2 和 IBM granite-4.0-h-tiny。我们选择这些模型是因为它们同时被 vLLM V0 和 V1 支持，而最近添加的模型（如 Qwen3-Next 和 Kimi Linear）仅在 V1 中支持。请参阅图 4 获取 V1 支持的完整混合模型列表。
结果：NVIDIA Nemotron-Nano-12B-v2
对于这种密集混合模型，V1 在 TTFT 和 ITL 方面总体上优于 V0，尽管两个版本之间不同的调度动态有时会导致权衡。吞吐量的提升是一致的，除了在极低并发下，此时 Mamba 内核中的 CPU 开销导致采用 PIECEWISE 的 V1 实际上表现得比 V0 更差。然而，我们可以看到 V1 采用 FULL_AND_PIECEWISE 解决了这个问题。总体而言，吞吐量提升在 2% 到 18% 之间。
[图 13：Nemotron-Nano-12B-v2 基准测试结果，对比 V0 和带有 PIECEWISE 和 FULL_AND_PIECEWISE CUDA Graphs 的 V1。]
结果：granite-4.0-h-tiny
granite-4.0-h-tiny 是一个具有 1B 激活参数的 7B MoE 模型。此类模型可以在相对较低的延迟下运行，这使得 CUDA Graphs 变得更加关键。在低并发下，V1 PIECEWISE 明显劣于 V0，但 V1 FULL_AND_PIECEWISE 持续带来了巨大的吞吐量提升，在某些情况下高达 91%，同时也降低了 TTFT 和 ITL 的延迟。
[图 14：granite-4.0-h-tiny 基准测试结果，显示使用 FULL_AND_PIECEWISE 时吞吐量大幅提升]
结论
混合模型不再是边缘案例。借助 vLLM V1，它们已完全集成到推理引擎中，受益于统一内存分配、CUDA Graph 优化以及与前缀缓存和 KV 传输等高级功能的兼容性。
对于开发者来说，结论很明确：混合架构可以扩展到更长的上下文而不会耗尽内存，通过将它们迁移到 V1，我们能够降低预填充和解码延迟，并实现显著的吞吐量增长，特别是对于具有 MoE 层的模型。
混合模型现在是构建企业级 AI 系统的实用工具。通过将混合模型提升为一等公民，vLLM 确保了开源社区能够跟上下一代大规模、长上下文工作负载的步伐。