FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling

发表时间: 2026-03 · arXiv:2603.06199

原文: https://arxiv.org/abs/2603.06199

作者/机构: Qihang Fan, Huaibo Huang, Zhiying Wu, Juqiu Wang, Bingning Wang, Ran He (MAIS&NLPR, CASIA; UCAS; WeChat, Tencent)

速读

一句话结论 本文提出了一种专为大型语言模型长上下文预填充阶段设计的加速框架 FlashPrefill，通过瞬时模式发现与基于最大值的动态阈值化技术，在过滤计算冗余的同时避免了排序开销，在 256K 上下文下实现了 27.78 倍的算子加速。

要解决什么问题 大型语言模型处理长上下文时，自注意力机制的二次方复杂度是核心卡点，在计算密集的预填充阶段尤为致命。现有稀疏注意力方法试图挖掘稀疏模式来降本，但存在两大瓶颈。首先是模式搜索延迟高：现有方法要么依赖复杂搜索策略引入沉重分析开销，要么采用无搜索方案但需维护密集的分数矩阵，导致极高的内存访问成本。其次是阈值筛选导致稀疏性不足与计算冗余：主流方法普遍采用 Top-k 或 Top-p 策略挑选注意力块，这在 GPU 上需要高昂的局部排序或前缀累加操作。更关键的是，注意力分数分布常呈重尾特征，Top-k/p 机制为了强行满足固定比例约束，会被迫保留大量毫无重要性的长尾注意力块，导致剪枝不彻底。此外，现有块稀疏核函数（如 Block-Sparse-Attention 库）普遍采用逻辑跳过策略，在内循环中依然遍历所有键块，仅用条件分支绕过掩码区域，导致线程仍需解码控制逻辑，引发严重指令流开销，在长序列下反而成为吞吐量瓶颈。

怎么做的 FlashPrefill 的核心思路是通过极低成本的近似计算快速定位高权重区域，用无排序的动态阈值截断低权重区域，并配合指令流优化的核函数完成稀疏计算。框架由三个关键部件构成。首先是瞬时模式发现机制，用于快速识别垂直、斜线和块状这三种常见稀疏模式。为避免长序列上直接计算查询与键的交互，该机制利用模型嵌入空间局部特征高度相似的特性，将块内所有键平均池化，用池化键 $\bar{k}$ 作为整个块的代理。数学上，该探测分数 $\Psi_{\mathrm{pool}}$ 是真实注意力贡献算术平均值的几何平均值下界： $$ \Psi_{\mathrm{pool}} = \exp (q \cdot \bar{k}) = \mathrm{GM}(e^{x_1}, \ldots, e^{x_n}) $$ 由于局部块内注意力方差极低，几何平均值可作为严格单调代理，在不改变块间排序的前提下大幅降低计算量。其次是用于核函数优化的块近似与融合归约。为消除探测过程中巨大中间矩阵的内存读写开销，作者开发了融合的 2D-Reduction 核函数。它不显式输出完整分数矩阵，而是在查询分块和池化键块间直接计算交互，并在查询轴上执行最大值归约和指数化，将细粒度分数压缩为 GPU SRAM 内每个块对的近似标量： $$ S_{I,J} = \sum_{q_i \in \text{Tile} I} \exp(q_i \cdot k_J - m_{I,J}) $$ 随后通过全局最大值重缩放进行概率映射，将内存占用从二次方降至块数量的平方级别。第三是基于最大值的动态阈值化机制。对于第 $I$ 个查询块，直接在候选键块中找到峰值分数，并乘以缩放因子 $\alpha$ 作为剪枝阈值： $$ \text{thresh}_I = \alpha \cdot \max_{J \le I}(\text{Score}_{I,J}) $$ 低于该阈值的键块直接丢弃。此设计只需单遍最大值归约，避开排序开销，且能自适应切断重尾分布中的冗余块。最后，底层引入索引驱动的物理跳转机制，直接将内存指针重定向到显著块坐标，彻底消除逻辑跳过策略的控制流停顿。

效果如何 实验在 NVIDIA H20 GPU 上搭建，评估了 Llama-3.1-8B-Instruct、Qwen3-30B-A3B-Instruct-2507 等语言及多模态模型，测试涵盖 RULER、InfiniteBench 等基准。对比基线包括精确计算路线的 Full Attention（基于 FlashAttention 2.8.3），以及代表不同稀疏路线的 MInference（动态稀疏）、FlexPrefill（上下文感知稀疏）、XAttention（反对角线评分块稀疏）和 FlashMoBA（混合块注意力）。量化结果显示，在 RULER 基准 128K 上下文下，FlashPrefill 在三个代表性模型上分别达到 22.67 倍、16.87 倍和 18.67 倍的算子加速；在 256K 长度下算子加速比高达 27.78 倍；短至 4K 上下文仍保持 1.71 倍加速。集成到 vLLM 框架后，128K 序列下的端到端首个令牌生成时间相比全注意力实现了 5.02 倍加速。该方法的代价是需针对不同模型在 4K 序列下校准唯一的超参数 $\alpha$（调节至计算密度约 70%）。此外，为保证模型性能，机制在极端稀疏时仍需硬性保留固定大小的注意力池（256 个令牌）和局部窗口（512 个令牌）不参与剪枝。

A1 主要贡献

核心问题: 长上下文建模是大型语言模型（LLM）的关键能力，但自注意力机制的二次方复杂度仍然是关键瓶颈，尤其是在计算密集的预填充（prefilling）阶段。

研究目标: 现有稀疏注意力方法通常存在显著的搜索延迟或稀疏性不足的问题。例如，Top-k或Top-p等选择策略需要在GPU上进行高开销的排序或累加操作，并且难以有效修剪影响力较小的token的长尾分布，导致计算冗余。本文旨在提出一个框架，通过瞬时模式发现和阈值化技术，实现超快的长上下文预填充，同时解决现有方法的效率和稀疏性问题。

创新点:
本文提出了FlashPrefill，一个为长上下文场景定制的超快预填充加速框架。其核心贡献如下：
1. 瞬时模式发现（Instantaneous Pattern Discovery）: 提出了一种瞬时模式发现方法，能够快速、准确地识别常见的注意力结构，如垂直、斜线和块状模式。为了加速这一过程，引入了块近似（block approximation）策略，通过简化内存访问和并行化发现过程，将模式发现的开销降至可忽略不计的水平。
2. 基于最大值的动态阈值化（Max-based Dynamic Thresholding）: 提出了一种基于最大值的动态阈值化机制，取代了传统的Top-k（需要局部排序）或Top-p（需要累加求和）策略。这种方法不仅避免了高昂的排序和累加开销，还能自适应地解决由注意力分数重尾分布引起的“不完全稀疏”问题，通过有效过滤冗余块实现更彻底的稀疏表示。
3. FlashPrefill框架: 结合上述策略，构建了FlashPrefill框架。在各种LLM和VLM上的广泛评估表明，该框架显著提升了预填充效率。例如，在256K序列长度上实现了前所未有的27.78倍加速，并且在短至4K的上下文长度上仍能保持1.71倍的加速，展示了其在不同序列尺度下的鲁棒性和实用性。


图 1 | 使用FlashPrefill在Qwen3-30B-A3B-Instruct-2507上进行的“大海捞针”评估，上下文长度从2K到256K。


图 2 | 在vLLM框架内，Qwen3-30B-A3B-Instruct-2507上相对于全注意力的端到端首个令牌生成时间（TTFT）加速比。


图 3 | 在Qwen3-30B-A3B-Instruct-2507上，各种算子相对于Flash Attention【5, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023, arXiv】的加速比。FlashPrefill展现出主导优势，尤其是在长上下文场景中。


图 4 | 不同方法中各个部分的执行时间。所有结果均在Qwen3-30B-A3B-Instruct-2507上测量。

A2 方法细节

3.1. 瞬时模式发现

现有方法的局限性: 根据先前研究【15, MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention, 2024, NeurIPS】【16, Flexprefill: A context-aware sparse attention mechanism for efficient long-sequence inference, 2025, ICLR】【27, Xattention: Block sparse attention with antidiagonal scoring, 2025, ICLR】，LLMs中存在三种主要的稀疏模式：垂直、斜线和块状稀疏。大多数现有方法依赖于基于搜索的策略来确定每个注意力头的特定模式，这通常会引入额外的计算分析开销【15, MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention, 2024, NeurIPS】【16, Flexprefill: A context-aware sparse attention mechanism for efficient long-sequence inference, 2025, ICLR】。而无搜索的替代方案通常需要更密集的计算并维护密集的注意力分数矩阵，导致巨大的计算和内存访问成本【27, Xattention: Block sparse attention with antidiagonal scoring, 2025, ICLR】。

基于结构特性的统一发现策略: 为了快速识别注意力图中的稀疏模式，我们对各种稀疏结构进行了定性分析。如图5所示，一个骨干的均匀分布查询集足以同时解决垂直、斜线和块稀疏模式。我们将这种方法的有效性归因于以下三个结构特性：
* 垂直模式（列不变性）: 这些结构代表全局键的显著性，其中特定的“锚点”令牌无论查询位置如何都能吸引大量注意力。由于这些特征在列方向上是不变的，稀疏的探测网格足以定位这些高能量列。
* 斜线模式（平移对称性）: 对角线图案通常源于局部句法依赖和相对位置偏差。由于它们在序列上的平移对称性，均匀分布的探针可以有效地采样这些局部“切片”，并解析出连续对角线结构的存在。
* 块稀疏模式（空间连续性）: 这些模式的特点是局部化的能量簇。利用这些区域的空间连续性，均匀采样策略确保了与这些密集簇相交的高概率，从而可以从稀疏的“命中”中稳健地推断出整个块的重要性。


图 5 | 模式发现示意图。红色虚线代表均匀分布的查询。

精确识别: 基于这些见解，我们直接使用完整的查询集来计算与所有键的注意力分数。为了确保最大程度的准确性，我们选择这种详尽的方法来精确识别显著块。

使用平均池化键进行优化: 直接在长序列上进行查询-键交互的计算成本过高。为了优化，我们采用平均池化的键 $\bar{k} = \frac{1}{n} \sum_{k_i \in B} k_i$ 作为块级别的代理。该方法基于LLM嵌入空间中固有的语义局部性和局部一致性，即局部块内的令牌表现出高的特征相似性和冗余的注意力模式【4, Longformer: The long-document transformer, 2020, arXiv】【26, Optimizing mixture of block attention, 2025, arXiv】【39, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023, NeurIPS】。

数学证明: 数学上，让一个块B包含n个键 $\{k_1, . . . , k_n\}$，其logit为 $x_i = q \cdot k_i$。我们定义探测分数 $\Psi_{pool}$ 和真实贡献 $\Psi_{sum}$ 如下：

$$\begin{aligned} \begin{aligned} \Psi_{\mathrm{pool}} & =\exp (q \cdot \bar{k})=\exp \left(\frac{1}{n} \sum_{i=1}^n x_i\right) \\ & =\left(\prod_{i=1}^n e^{x_i}\right)^{1 / n}=\mathrm{GM}\left(e^{x_1}, \ldots, e^{x_n}\right) \\ \Psi_{\text {sum }} & =\sum_{i=1}^n e^{x_i}=n \cdot \operatorname{AM}\left(e^{x_1}, \ldots, e^{x_n}\right) \end{aligned} \end{aligned}$$

其中GM(·)和AM(·)分别表示几何平均值和算术平均值。根据算术-几何平均值不等式（AM-GM Inequality）：

$$\frac{1}{n} \Psi_{\text{sum}} \geqslant \Psi_{\text{pool}}$$

虽然 $\Psi_{pool}$ 是一个下界，但由于注意力分布特有的低块内方差（$\sigma^2 \rightarrow 0$），块间的排序不变性得以维持。在这种情况下，几何平均值作为算术平均值的严格单调代理，确保了块的相对排序保持不变，且秩失真可忽略不计。在计算单个查询的注意力分数后，我们对每个块内的所有查询执行二次平均操作，以得出聚合的块级别重要性。

3.2. 用于核函数优化的块近似

问题: 上述方法在计算块内注意力平均值时，仍然需要存储和计算一个巨大的 $N \times (N/B)$ 中间矩阵，其中B是块大小。为了解决这个问题，我们再次利用同一块内令牌的语义相似性【4, Longformer: The long-document transformer, 2020, arXiv】【26, Optimizing mixture of block attention, 2025, arXiv】【39, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023, NeurIPS】。鉴于局部块中的令牌表现出显著的一致性，它们产生的注意力分布高度冗余。因此，我们放弃了计算成本高昂的严格注意力计算，转而采用块级注意力分数近似。具体来说，我们实现了一个融合的2D-Reduction核函数，将计算从显式的“计算后池化”序列转变为单遍融合核函数，显著减少了全局内存流量。图6比较了这两种方法。


图 6 | 块评分方法比较。（左）标准核函数，显式物化注意力分数。（右）我们优化的核函数，利用块级注意力近似。与标准实现相比，我们提出的块级注意力分数近似通过绕过 $O(N^2/B)$ 的中间内存流量，显著降低了内存访问开销。这里，N表示序列长度，B表示块大小，d表示令牌维度。

融合的块级注意力近似: 为了避免严格的令牌级计算带来的高昂成本，我们放弃精确的注意力计算，转而采用块级近似。通过将查询tile大小设置为等于块大小(B)，核函数计算查询tile和池化键块之间的交互。我们利用查询维度上的1D-reduction将细粒度分数压缩为GPU SRAM内每个块对的单个近似标量：
* 分块交互 (Tiled Interaction): 对于每个查询tile，核函数加载池化的键块（其中块中的每个键由单个向量 $\bar{k}_J$ 表示），并计算交互 $q_i^T \bar{k}_J$ 作为代表性代理。
* 稳定的在线归约 (Stable Online Reduction): 为了在近似Softmax行为的同时保持数值稳定性，我们在查询轴上执行最大值归约和指数化：

$$m_{I,J} = \max_{q_i \in \text{Tile}I} (q_i \cdot k_J)$$

$$S_{I,J} = \sum_{q_i \in \text{Tile} I} \exp(q_i \cdot k_J - m_{I,J})$$

核函数在查询tile维度上执行求和，为每个块对(I, J)生成近似的能量分数 $S_{I,J}$ 和局部最大值 $m_{I,J}$。

保持一致性的全局归一化: 虽然块级近似 $S_{I,J}$ 使用局部最大值进行了稳定化处理，但这些分数必须对齐以确保在整个序列中的全局可比性。我们执行一个全局归一化过程，将这些近似值转换为统一的重要性图：
* 最大值重缩放 (Maximum Rescaling): 我们识别每个查询tile在所有块中的全局最大值，$M_I = \max_J (m_{I,J})$，并将局部近似分数重缩放到一个共同的分母上：

$$\mathcal{S}_{I, J}^{\prime}=\mathcal{S}_{I, J} \times \exp \left(m_{I, J}-\mathcal{M}_{I}\right)$$

这个融合过程确保了最终的块级图能准确反映相对注意力分布，同时将内存占用从 $O(N \cdot N/B)$ 减少到 $O((N/B)^2)$。通过放弃严格计算转而采用这种高保真度的近似，我们实现了即使对于超长序列也能进行近乎瞬时的模式发现。

3.3. 与先前方法的比较

方法对比: 与先前的模式发现方法相比，我们提出的方法更快、更准确。基于Llama-3.1-8B-Instruct，我们比较了三种不同的模式发现方法：1) 对查询和键块都应用均值池化，然后计算块间注意力分数【15, MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention, 2024, NeurIPS】【16, Flexprefill: A context-aware sparse attention mechanism for efficient long-sequence inference, 2025, ICLR】；2) 使用3.1节中描述的原始方法，不进行3.2节中描述的任何块近似【19, Moba: Mixture of block attention for long-context llms, 2025, arXiv】【26, Optimizing mixture of block attention, 2025, arXiv】；以及3) 我们提出的基于块近似的方法。

结果分析: 结果如表1所示。我们在RULER【14, Ruler: What’s the real context size of your long-context language models?, 2024, arXiv】基准上进行了评估；为确保不同模式发现方法之间的公平比较，我们始终选择得分最高的8个块。方法1)由于其过于激进的估计策略而导致显著的性能下降。相比之下，方法2)由于过多的内存访问成本而产生巨大的时间开销。我们提出的方法3)在效率和效果之间达到了最佳平衡。


表 1 | 不同模式发现方法的比较。

3.4. 基于最大值的动态阈值化

现有方法的效率问题: 先前的方法依赖Top-k或Top-p启发式算法来选择显著的令牌或块【15, MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention, 2024, NeurIPS】【16, Flexprefill: A context-aware sparse attention mechanism for efficient long-sequence inference, 2025, ICLR】【19, Moba: Mixture of block attention for long-context llms, 2025, arXiv】【26, Optimizing mixture of block attention, 2025, arXiv】【27, Xattention: Block sparse attention with antidiagonal scoring, 2025, ICLR】。然而，这些方法需要相对昂贵的排序和累积计算，在效率方面并非最优。


图 7 | 不同阈值化方法的比较。（左）Top-k和Top-p选择策略。（右）我们提出的基于最大值的阈值化。Top-k和Top-p方法通常易受长尾分布的影响，因为它们可能仅仅为了满足固定的k或p约束而包含大量低重要性的块。相比之下，我们的动态阈值化有效地减轻了长尾分布的影响，通过更精确地选择显著块实现了更高的稀疏度。

现有方法的稀疏性问题: 此外，Top-k和Top-p机制由于对长尾分布的高度敏感性而导致稀疏性不足。如图7所示，当注意力分数分布呈长尾状时（这在大多数情况下很普遍），这些方法通常需要包含大量不重要的块，仅仅是为了满足固定的k或p约束，从而未能实现最佳的剪枝效率。

提出的动态阈值化机制: 为了规避这些限制，我们引入了一种基于最大值的动态阈值化机制。具体来说，对于第I个查询块，我们在所有候选键块中识别出峰值注意力分数，并直接从这个最大值推导出剪枝阈值：

$$\text{thresh}_I = \alpha \cdot \max_{J \le I}(\text{Score}_{I,J})$$

其中 $\alpha$ 是一个可调的缩放因子。对于查询块I，任何分数低于 $thresh_I$ 的键块J都将被从计算中丢弃。至关重要的是，这种方法只需要单遍最大值归约，而不是计算成本高昂的全局排序，从而显著提高了执行效率。此外，如图7所示，这种动态阈值化有效地减轻了长尾分布的干扰，通过专注于真正显著的块，与传统启发式方法相比实现了更优的稀疏性。

表 2 | 块稀疏注意力实现的延迟比较。

3.5. 优化的块稀疏注意力核函数

背景: 在获得查询块和键块之间的稀疏模式后，FlashPrefill对这些块执行标准的块稀疏注意力。在之前的工作中【25, Proxyattn: Guided sparse attention via representative heads, 2025, arXiv】【27, Xattention: Block sparse attention with antidiagonal scoring, 2025, ICLR】，这种计算主要基于Block-Sparse-Attention代码库【13, Block Sparse Attention, 2024, GitHub】实现。

现有实现的瓶颈: 我们观察到，Block-Sparse-Attention仓库中的实现未能充分利用稀疏执行的潜在效率。该仓库对被掩码的块采用逻辑跳过策略：其内循环不加选择地遍历所有键块的线性范围，使用条件分支绕过被掩码区域的GEMM操作。即使一个块被掩码且其矩阵乘法被绕过，线程执行流仍必须获取、解码和执行该迭代内的循环控制逻辑、指针算术和同步原语。这种冗余的指令流开销导致了显著的流水线利用率不足，随着序列长度的增加，实际上成为了核函数整体吞吐量的瓶颈。

优化方案: 为了进一步提升块稀疏注意力的性能，我们优化了其执行模型。具体来说，我们放弃了受指令流开销影响的逻辑跳过策略，并实现了一种索引驱动的物理跳转机制。通过直接将内存指针重定向到显著块的坐标，我们的方法消除了冗余的控制流处理和同步停顿，从而在长序列场景中最大化硬件吞吐量和计算强度。表2对两种方法在不同密度和序列长度下的效率进行了基准测试，我们的实现显著优于现有的基线。

A4 实验环境

A4 实验结果

InfiniteBench: 如表3所示，FlashPrefill在InfiniteBench基准上，无论是在密集模型还是混合专家（MoE）模型上，都一致性地取得了优于其他基线的性能。


表 3 | 不同方法在InfiniteBench上各种模型和任务的性能比较。

RULER: 如表4所示，FlashPrefill在RULER基准上的所有测试长度上都实现了显著的加速。具体来说，在128K上下文长度下，FlashPrefill在三个代表性模型上分别达到了22.67倍、16.87倍和18.67倍的加速，远超现有方法。


表 4 | 不同模型和方法的性能与效率对比。报告了评估分数（左）和相对于全注意力的算子加速比（右）。

VideoMME: 表5展示了在Video-MME基准上的性能。与现有的稀疏注意力方法相比，FlashPrefill取得了更优异的结果。


表 5 | 不同方法在VideoMME上各种模型的性能比较。

密度: 表6展示了不同方法在Qwen3-30B-A3B-Instruct-2507模型上的密度。可以观察到，FlashPrefill有效减轻了长尾分布的影响。随着序列长度增加，有效信息相对减少，因此与其他两种方法相比，FlashPrefill的密度显著降低。


表 6 | Qwen3-30B-A3B-Instruct-2507上各种方法的密度。

端到端首个令牌生成时间（TTFT）加速比: 如表7所示，将FlashPrefill集成到vLLM框架后，在包括密集和MoE架构在内的三种语言模型上，相比全注意力实现了显著的性能提升。特别是在Qwen3-30B-A3B-Instruct-2507上，在128K序列长度下实现了5.02倍的加速。


表 7 | FlashPrefill实现的端到端TTFT加速比。

消融研究:


表 8 | 不同阈值化策略的性能得分和注意力密度。

A5 结论

本文提出了FlashPrefill，一种旨在显著加速大型语言模型（LLMs）长上下文预填充阶段的新方法。FlashPrefill引入了一种瞬时模式发现机制，并通过基于块近似的核函数实现进一步优化，以最小化内存访问开销。此外，我们提出了基于最大值的动态阈值化方法，该方法无需排序和累积操作，同时减轻了长尾分布的影响，从而实现了更高的稀疏度。我们在各种LLMs、VLMs和多样化的基准上进行的广泛评估表明，FlashPrefill在保持卓越性能的同时，有效提升了预填充效率。

A6 附录

A. 超参数配置

唯一的超参数 $\alpha$: 在FlashPrefill中，唯一需要调整的超参数是用于确定动态阈值的 $\alpha$ 值。我们通过观察不同模型在4K序列长度下的计算密度来校准这个阈值。具体来说，我们调节 $\alpha$ 以使4K序列的计算密度维持在约70%，如表6所示。

固定参数: 我们对所有模型统一使用128的块大小。除了通过基于最大值的动态阈值化选择的块之外，我们还明确保留了注意力池（attention sinks）和局部窗口，其大小分别设置为256和512个令牌。


表 9 | 各种模型的超参数配置及产生的计算密度。

B. 详细实现

框架分解: 为了更清晰地阐述我们提出的FlashPrefill，我们将该框架分解为三个不同阶段：(i) 瞬时模式发现，(ii) 基于最大值的动态阈值化和块选择，以及 (iii) 块稀疏注意力。每个阶段的详细逻辑分别在算法1、算法2和算法3中以伪代码形式正式化。


算法 1 | 瞬时模式发现


算法 2 | 基于最大值的动态阈值化和块选择


算法 3 | 块稀疏注意力