FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling
微信 · 中科院自动化所 | arXiv:2603.06199 | GitHub: qhfan/FlashPrefill
文本版 · 供搜索与朗读
FlashPrefill Poster
FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling
微信 · 中科院自动化所 | arXiv:2603.06199 | GitHub: qhfan/FlashPrefill
💡 核心突破：“寻找重要块”和“筛选策略”做到近乎零成本 + 动态阈值剪掉注意力长尾 + 索引压紧实现物理跳转
🚀 在 256K 上下文下实现 27.78倍 算子加速
背景与动机
Prefill 阶段瓶颈：
Prefill 需遍历 Prompt 计算 KV Cache，复杂度 $O(L^2)$。
长上下文导致 TTFT (首字延迟) 爆炸式增长。
用户界面卡顿，体验极差。
Prefill 耗时随长度指数增长
现有方法痛点
❌ 模式发现延迟高
估算 Block 重要性本身计算量大。
❌ 筛选策略昂贵
Top-k 排序、Top-p 累加难以并行。
❌ 稀疏不彻底
长尾分布下，为凑够 K 个/P 概率引入大量无效块。
FlashPrefill 方法
1. 瞬时模式发现
近似运算
Block 级近似，均值代理排序。
低方差假设：块内语义相似。
显存优化
Key/Query 池化 + 全局再加权。
生成全局“注意力地图”。
📉 数据搬运从 $L imes L/B$ 降至 $(L/B)^2$。
2. 动态阈值剪枝
Threshold = α × max(Score)
计算极轻：仅需一次 Max-reduction。
剪掉长尾：不凑数，彻底丢弃低分块。
🛡️ 安全网机制：
保留 Attention Sinks (前256)
保留 Local Window (近512)
3. 索引压紧物理跳转
问题： 逻辑跳过 (if mask=0) 仍遍历所有块，分支开销大。
方案：
压紧有效块索引为连续列表。
Kernel 内层仅遍历有效索引。
逻辑跳过 → 物理跳转
内存访问更集中，循环次数骤减
实验结果
⚡ 性能加速
27.78x
Prefill 算子加速 (Qwen3-30B @ 256K)
稀疏度对比 (256K):
FlashPrefill: 3.5% (保留块)
FlexPrefill: 8.4%
XAttention: 18.5%
🎯 准确率保持
RULER (Qwen3-30B): 92.68 vs Full 93.28 (微降)
InfiniteBench (Qwen2.5-7B): 24.93 vs Full 23.87 (略优)
Needle-in-Haystack: 2K~256K 检索能力几乎无损。
端到端 TTFT 显著降低 (集成 vLLM)
总结：FlashPrefill 通过近似计算、动态阈值和物理跳转，实现了长上下文 Prefill 的零成本稀疏化，极大提升了推理效率。