FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling
💡 核心突破：“寻找重要块”和“筛选策略”做到近乎零成本 + 动态阈值剪掉注意力长尾 + 索引压紧实现物理跳转
🚀 在 256K 上下文下实现 27.78倍 算子加速
🚀 在 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)