模型链接：https://huggingface.co/openbmb/MiniCPM-Llama3-V-2_5
论文链接：https://arxiv.org/abs/2404.06395

引言：破解长序列处理瓶颈，InfLLM-V2 框架发布

随着大型语言模型（LLM）的快速发展，其在长序列处理上的能力已成为衡量其应用价值的关键指标。无论是处理深度研究资料、赋予聊天机器人长期记忆，还是解决复杂的软件工程问题，模型理解与生成长序列的能力都直接决定了其在真实场景中的表现。然而，标准 Transformer 架构中的自注意力机制在处理长序列时，其计算与内存开销会随序列长度呈二次方增长，构成了严峻的瓶颈。

为应对此挑战，可训练稀疏注意力机制成为一个极具潜力的研究方向。然而，现有的代表性方法（如 NSA 架构）与主流的“短序列预训练-长序列微调”范式存在显著的架构错配问题，不仅引入了过多的额外参数，还导致模型在迁移训练中收敛不稳，并为短序列处理带来不必要的开销。

针对这些问题，清华大学与 OpenBMB 联合提出了 InfLLM-V2：一个零额外参数、可实现稠密与稀疏模式无缝切换的原生注意力框架。它在短文本场景下保持与原生稠密注意力一致的高性能，而在长文本场景下则平滑过渡到高效的稀疏模式，实现了端到端的显著加速。

---

核心创新：InfLLM-V2 如何实现“强”与“快”的统一

InfLLM-V2 建立在免训练稀疏注意力机制 InfLLM 的基础上，通过三大核心创新，解决了现有可训练稀疏方法的关键痛点：

1. 无缝的短长序列自适应 (Seamless Short-to-Long Adaptation)：
   不同于 NSA 引入三套独立键值（KV）投影参数与多个注意力模块的复杂设计，InfLLM-V2 通过无参数的架构调整，直接复用已有的稠密注意力参数。这一设计消除了架构层面的不匹配，完美对齐“短序列预训练-长序列微调”范式，确保了训练过程的稳定性与快速收敛。

2. 长短序列双重效率优化 (Efficiency for Both Short and Long Sequences)：
   由于零新增参数且分布偏移极小，InfLLM-V2 在经过长文本微调后，依然可以在处理短序列时无损地切换回原生稠密注意力模式，避免了任何性能衰退或额外计算开销。对于长序列，则采用统一的稀疏注意力范式，在 prefill 和 decode 阶段实现全链路加速。

3. 硬件友好的高效块选择机制 (Accelerated Block Selection Mechanism)：
   稀疏注意力的核心在于“块选择”（Block Selection）步骤，即为每个查询（Query）确定最相关的上下文块。然而，这一步骤本身可能成为新的效率瓶颈。InfLLM-V2 提出了一套硬件感知的优化实现，通过创新的计算核融合技术，显著降低了高带宽内存（HBM）的I/O与计算负载，从而完全释放稀疏注意力的性能潜力。

---

技术解析：InfLLM-V2 的框架设计

标准自注意力机制中，每个查询词元 (Query, Q) 都需要与历史序列中的所有键值词元 (Key-Value, K-V) 进行交互。稀疏注意力的目标是将其优化为每个查询词元仅与一个精心挑选的上下文子集交互，主要包含两个阶段：块选择与稀疏注意力计算。

对齐计算：架构的简化与统一

InfLLM-V2 的设计哲学是简洁与对齐。

• 共享键值投影 (Shared Key-Value Projection)：框架摒弃了为不同注意力模式设置多套KV投影矩阵的做法，转而使用单一共享的 W_K 和 W_V 矩阵。在长文本微调阶段，这些矩阵直接从预训练好的稠密模型中继承，保证了参数的连续性。

• 统一的注意力计算 (Aligned Computation)：NSA 架构包含三种注意力模块（压缩、选择、滑动）和一个额外的门控模块来聚合输出。这种设计在短序列上会产生巨大开销。InfLLM-V2 将此多分支结构合并为单一分支：

  • 融合稀疏模式：将 NSA 中的“选择注意力”（Selected Attention）与“滑动窗口注意力”（Sliding Attention）所关注的区域取并集，形成一个统一的稀疏注意力模式。通过扩展局部块的数量 N_local，使其严格覆盖滑动窗口区域 w，即满足 N_local ≥ ⌈w/B⌉ + 1，其中 B 为块大小。
  • 简化压缩注意力：取消了“压缩注意力”（Compressed Attention）的输出路径，仅保留其生成的注意力分数 S_cmp 用于后续的块选择，这使得整体架构更接近标准的单输出注意力形式。

高效块选择：从多阶段压缩到硬件优化

为了在块选择阶段既保留关键信息又降低计算成本，InfLLM-V2 进行了两方面的优化。

1. 无参数的多阶段池化：
   将原先基于 MLP 的有参数压缩操作替换为无参数的池化操作。具体采用一个三阶段、由粗到细的压缩流程：首先通过均值池化（Mean-Pooling）生成一个粗粒度的键序列 K_C1，然后计算初步注意力分数 S_C1。
   $$K_{C{1}_i}=\text{Mean}(K_{i\cdot s_{C1}:i\cdot s_{C1}+l_{C1}})$$
   $$S_{C1}=\text{Softmax}(Q(K_{C1}{)}^{\top })$$
   随后，在组查询注意力（GQA）的头组（head group）内对分数进行求和，并在最终阶段使用最大池化（Max-Pooling）保留最显著的特征，得到最终用于 Top-K 选择的压缩分数 S_cmp。

2. LSE 近似与计算核融合：
   我们发现，计算压缩分数 S_cmp 的过程，尤其是中间结果 S_C1 的存储，会带来巨大的 HBM I/O 负担，成为性能瓶颈。借鉴 FlashAttention 的思想，我们设计了**融合头组求和（Fused Head Group Summation）的高效计算核。为了解决在线 Softmax 与头组求和的计算顺序冲突，我们采用了一种两遍（two-pass）方法：第一遍使用一个更粗粒度的键表示 K_C2 来近似计算 Softmax 的归一化项（log-sum-exp, lse），第二遍则利用这个 lse 来计算精确的注意力分数、在 SRAM 内完成头组求和，并最终只将聚合后的分数写回 HBM。这种LSE 近似（LSE Approximation）**方法将计算开销从 2 倍降低到 1.25 倍，同时大幅减少了 I/O。

---

实验验证

我们在长文本理解和长链推理任务上对 InfLLM-V2 进行了全面评估。实验证明，该框架在大幅提升效率的同时，能够高度保留稠密模型的性能。

长文本理解任务

在 RULER、LongBench 和 LongPPL 等权威长文本评测基准上，InfLLM-V2 的性能与全量微调的稠密注意力模型（FULLATTN）几乎持平，并显著优于其他所有稀疏注意力方法。相比之下，NSA 方法由于引入大量新参数，在有限的长文本数据微调后，难以有效捕捉长距离依赖关系，导致性能下降明显。

长链推理任务

在数学和代码等需要深度思考和长链推理的任务中，模型的加速能力尤为重要。实验结果表明，InfLLM-V2 在保持与稠密模型相当性能的同时，展现了稀疏注意力在加速复杂推理过程中的巨大潜力。

效率评测

我们在 NVIDIA A100 和 RTX 4090 两款 GPU 上进行了推理效率测试。

• 算子层面：在 128K 的长下文中，相较于 FlashAttention-2 实现的稠密注意力，InfLLM-V2 实现了 4 到 9 倍的算子层面加速。消融实验证明，我们提出的高效块选择设计是关键的加速来源。
• 端到端层面：在 prefill（首次生成 token）和 decode（后续 token 生成）阶段，InfLLM-V2 分别实现了约 2.1 倍和 2.3 倍的端到端加速。

---

开源实践：首个原生稀疏注意力模型 MiniCPM-V2

基于 InfLLM-V2 框架，OpenBMB 与清华大学已训练并开源了 MiniCPM-V2 系列模型，这是社区首个开源的原生稀疏注意力模型。该模型充分利用了稀疏注意力、推测采样等高效算法，在代码、数学等推理任务的测试中，其推理速度比同尺寸的 Qwen2、Llama3 等开源模型快 3 倍以上，为研究社区提供了一个可复现的高效长文本模型实现。