FlashPrefill V2：面向长上下文 LLM 服务的块稀疏 Prefill 注意力

均值校正+FA3/4级稀疏内核，长上下文prefill最高加速47倍并落地SGLang

前置知识

Prefill 与 Decode 阶段

LLM 推理分两阶段：prefill 一次性处理整段输入 prompt，注意力计算量为长度平方，决定首 token 延迟（TTFT）；decode 逐 token 生成，受显存带宽限制，决定每 token 延迟（TPOT）。长上下文下 prefill 注意力是主要瓶颈。

本文加速目标正是 prefill 注意力，理解两阶段差异才能明白为何 decode 退回稠密注意力、以及端到端收益为何受注意力占比约束。

FlashAttention 与 warp 专用化

FlashAttention 是 IO 感知的精确注意力内核，通过分块与 online softmax 避免 $O(L^2)$ 中间矩阵写显存。FA3/4 在 Hopper 上引入 TMA 异步搬运、warp 专用化流水线与 GEMM-softmax 重叠（pingpong）。

本文第 3.3 节把稀疏算子完全重写为 FA3/4 同级流水线，不懂这套执行模型就看不懂内核设计与公平基线的含义。

块稀疏注意力

将 K/V 切成固定大小块（本文 128 token），用低成本估计每块重要性，仅对少量关键块做精确注意力，其余跳过。本文用均匀分布的探测查询对块均值池化键打分，并以阈值 $\alpha\cdot\max$ 保留块，免去 Top-k/Top-p 的全局排序开销。

这是从 V1 继承的核心算法框架，V2 的均值校正、PackGQA 与内核改造全部围绕它展开。

Paged KV Cache 与连续批处理

vLLM 提出的分页管理把 KV 缓存切成固定页、经页表寻址以消除显存碎片；连续批处理允许请求动态进出当前批次以最大化吞吐。vLLM、SGLang 等生产引擎都依赖这两项技术。

多数研究型稀疏内核假设连续 KV 布局，无法进入生产系统；原生兼容二者正是本文三大贡献之一，也是与以往工作的关键分野。

FP8 量化

用 8 位浮点 e4m3 替代 BF16 存储与计算，Tensor Core 吞吐翻倍、带宽减半，但需缩放因子 $c_q,c_k,c_v$ 管理动态范围；softmax 概率需平移映射到 [0,256] 才能充分利用 e4m3 表示空间，偏移在 softmax 比值中相消。

FP8 变体是加速从 27× 翻到 47× 的关键，且量化压缩分数边际后均值校正的必要性显著放大，是消融实验的重要轴线。

GQA（分组查询注意力）

让多个查询头共享同一组 KV 头（共享比例 $g=H_q/H_{kv}$）以压缩 KV 缓存体积。内核层面 naive 映射会让 $g$ 个线程块重复装载同一 KV 块，浪费带宽与共享内存。

本文的 PackGQA 把查询重排为 $(gL)\times H_{kv}\times d$，一个线程块覆盖整组查询头，是对齐 FA3/4 与缩小索引元数据的关键设计。

研究动机

长上下文 LLM 的自注意力复杂度随序列长度平方增长，在计算密集的 prefill 阶段成为服务瓶颈：论文实测 FA3/4 稠密后端在 128K、batch 16 时 Qwen3-30B-A3B 的 TTFT 达 123.2 秒。稀疏注意力（MInference、FlexPrefill、XAttention）虽只计算关键块，但普遍停留在算法原型层面：其一，激进稀疏下被剪块的 softmax 概率质量不可忽略、精度退化不可控，消融显示无补偿的 FP8 变体在 128K RULER 上掉 6.2 分；其二，内核沿用 FlashAttention-2 的同步流水线，在 TMA 与 wgmma 主导效率的 Hopper GPU 上远落后于 FA3/4；其三，KV 布局假设连续存储，与生产引擎必需的 paged KV cache 和连续批处理不兼容，无法作为注意力后端集成。团队前作 FlashPrefill 用瞬时模式发现与 max 动态阈值消除了 Top-k 排序延迟，但同样面临这三重差距，距离生产部署尚远。

本文的目标是本文目标是把 FlashPrefill 从算法原型演进为可实际部署的长上下文 prefill 方案，沿三个维度补齐差距：(1) 精度可控——引入均值校正项，用被剪块的池化 K/V 统计补偿被丢弃的概率质量，使 128K 下密度低于 5% 的极端稀疏仍保持与全注意力 1.8 分以内的 RULER 差距；(2) 内核对齐——把稀疏算子重写为与 FlashAttention-3/4 完全同级的 Hopper 内核，包含 PackGQA 内存访问、warp 专用化 producer-consumer 流水线、pingpong GEMM-softmax 重叠，并扩展到 FP8 推理；(3) 系统就绪——原生支持 paged KV cache 与连续批处理，以标准注意力后端集成进 SGLang，不改动模型定义、KV 布局与调度逻辑。量化目标是在广泛部署的 H20 上取得数十倍算子加速，端到端 TTFT 显著下降，且平均精度损失控制在约 1 分内。

与已有工作不同的是，多数稀疏注意力研究把稀疏模式当作唯一设计轴，忽视部署三要素。本文的独特切入是可控精度、内核对齐、系统兼容三位一体：(1) 不再剪枝即丢弃——给被剪块一个零阶均值替代项 $|\mathcal{B}_J|e^{\bar{s}_J}\bar{v}_J$ 补偿分子分母中被丢弃的质量，并给出完整误差理论（质量项二阶精确、分子保留一阶块内协方差），显著拓宽可用稀疏范围；(2) 坚持用与稀疏算子共享流水线的 FA3/4 对齐稠密内核作基线，把稀疏算法收益与内核工程收益干净分离，避免虚高加速比；(3) 从第一天起围绕 paged KV、变长请求与连续批处理设计，CSR 索引直接兼容 SGLang 格式，并与生产级 HPC-Ops BSA 内核逐项对照（Tab.7），证明优势来自设计本身而非牺牲生产特性换速度。

核心方法

直觉：长上下文的注意力分布高度局部化，少数关键块承载绝大部分概率质量，但密度低于 5% 时直接剪块会丢失不可忽略的质量；且若 Hopper 上的内核效率跟不上，稀疏性无法兑换成实际加速。FlashPrefill V2 的方案是精确计算关键块、只用一次均值近似其余块，并配上 FA3/4 级内核。每层执行两阶段：Stage 1 索引阶段——PackGQA 打包查询，对块均值池化的键做单次融合 GEMM 打分（在线维护 tile 最大值 $M$ 与块能量 $E_J$），用 $\alpha\cdot\max$ 阈值选块并原地压缩成 CSR 索引，强制保留 sink/局部窗口/对角块；Stage 2 注意力阶段——warp 专用化内核遍历选中块做精确 wgmma 与 online softmax（pingpong 重叠），未选中块以均值向量形式作为带 $\log|\mathcal{B}_J|$ 平移的额外迭代注入同一流水线完成均值校正，两条流在 online softmax 内合并输出。

核心创新是均值校正项：被剪块不丢弃，而是用块均值 $(\bar{k}_J,\bar{v}_J)$ 以替代项 $|\mathcal{B}_J|e^{\bar{s}_J}\bar{v}_J$ 补入 softmax 分子、$|\mathcal{B}_J|e^{\bar{s}_J}$ 补入分母（$\bar{s}_J=q\cdot\bar{k}_J/\sqrt{d}$），恢复被丢弃的概率质量。理论上分母替代二阶精确，分子残留一阶块内协方差 $\delta_v$；max 阈值保证每个被剪块质量份额 $\lesssim\alpha$，相对直接丢弃的误差比 $\sim O(\delta_{\mathrm{rms}})$。与已有方法的本质区别：先前方法视稀疏模式为唯一设计轴、剪枝即丢弃；本文把补偿融合进内核主循环——作为 logit 平移 $\log|\mathcal{B}_J|$ 的额外迭代，复用 MMA 流水线、零额外内核启动——因此 FP8 压缩分数边际后依然有效，并把可用稀疏范围扩大到 $\alpha=0.2$（密度减半至 5.2% 仍得 80.12 分）。

方法步骤详情

(1) PackGQA 把查询矩阵 $Q$ 重排为 $(gL)\times H_{kv}\times d$，一个线程块覆盖整组查询头；(2) 一次 Triton pass 算出块均值 $(\bar{k}_J,\bar{v}_J)$；(3) 单次 GEMM 融合打分，块能量 $E_J$ bitcast 直接写入索引缓冲，省去独立分数缓冲与第二次 GEMM；(4) 阈值 $\hat{E}_J\ge\alpha\cdot\max\hat{E}$ 或属 sink/window/diag 者保留，原地压缩成 CSR 索引；(5) 持久调度器按因果负载降序派发，producer 经页表 cp.async/TMA 装载 K/V，consumer 按 pingpong 顺序执行 QK GEMM$(n)$、PV GEMM$(n{+}1)$、online softmax$(n)$；(6) 二指针扫描把未选中块聚合为 bitmask，作为平移 $\log|\mathcal{B}_J|$ 的校正迭代注入主流水线；(7) epilogue 按累计分母重缩放并反置换 FP8 列序，输出注意力结果 $O$。

技术新颖性

三重新颖性。(1) 算法层：首个把剪块补偿理论化并融合进内核主循环的 prefill 稀疏方案——式(7)-(11) 证明质量替代二阶精确、分子残留一阶协方差、相对丢弃的误差比 $\sim O(\delta_{\mathrm{rms}})$，且 max 阈值给出每块质量份额 $\lesssim\alpha$ 的关键界；实测在 FP8 128K 挽回 6.2 分。(2) 内核层：把 FA3/4 的 TMA、warp 专用化与 pingpong 重叠扩展到块稀疏场景——索引驱动遍历使 producer/consumer 控制流构造性一致、免跳块同步；PackGQA 消除 $g$ 倍重复 KV 装载并把索引元数据缩为 $1/g$；CSR 按选中块数均分实现稀疏感知负载均衡；单 pass 打分消除第二次 GEMM。(3) 系统层：研究型稀疏内核中少有的原生 paged KV、变长请求、连续批处理与 SGLang 后端集成，对照生产级 HPC-Ops BSA 仍快 6-7%（后者仅支持 FP8、用稠密掩码索引，64K 需 8MB，本文 CSR 仅 3MB）。

**Figure 3**: The FlashPrefill V2 prefill pipeline (Stage 1: Block Selection, Stage 2: Sparse Attention Kernel)

**Figure 4**: Mean correction: selected blocks are computed exactly, pruned blocks are pooled into mean statistics contributing surrogate terms

实验结果

(1) RULER（Tab.1）：与全注意力平均差 ≤1.1 分（87.79/86.23/91.76 vs 88.82/87.06/92.05），128K 差距仍 <1.8 分；128K 算子加速对 FA2 达 27.2×-27.5×（BF16）、47.3×-47.6×（FP8），对 FA3/4 稠密基线仍 17.5×/30.5×，远超基线（≤5.4×）及前作。(2) LongBench（Tab.2）：三模型均列稀疏第一（49.31/46.96/50.73 vs 最强基线 48.25/45.90/49.61）。(3) 端到端（Tab.4）：SGLang TP=4 下 128K TTFT 降 2.1×-3.4×（BF16）、3.7×-4.8×（FP8），Qwen3-30B batch16 123.2s→36.2s/25.5s。(4) 开环服务（Tab.5）：FA3/4 饱和于 0.31-0.37 req/s，V2 吞吐翻倍、FP8 达 0.88-1.34 req/s。(5) 消融（Tab.8/9）：无校正 FP8 128K 掉 6.2 分；校正后 α=0.2 密度减半仍 80.12 分。

**Table 1**: RULER scores (left) and attention operator speedups over FlashAttention-2 (right) at sequence lengths from 4K to 128K

**Table 2**: Performance comparison on all 21 tasks of LongBench

**Table 3**: Attention density of FlashPrefill V2 measured on needle-in-a-haystack inputs at each sequence length

**Table 4**: End-to-end time-to-first-token (TTFT, seconds) served by SGLang, batch sizes 1 to 16

**Table 5**: Open-loop serving results under Poisson arrivals with mixed prompt lengths (4K-128K), 100 requests per cell

**Table 6**: Open-loop serving at 16 req/s with chunked prefill at chunk sizes of 8K and 16K

**Table 7**: Design comparison with the HPC-Ops block-sparse attention kernel

**Table 8**: Ablation of the mean correction term on RULER with Qwen3-4B-Instruct-2507, in both BF16 and FP8

**Table 9**: Ablation of the selection threshold α on RULER at 64K with Qwen3-4B-Instruct-2507 in FP8

**Figure 2**: Speedup of various attention operators relative to FlashAttention-2 on NVIDIA H20 GPUs, batch size 4

**Figure 5**: Block-sparse attention latency at 64K sequence length in FP8; comparison with HPC-Ops BSA

查看结构化数据

任务	指标	本文	基线	提升
RULER 长上下文基准（Llama-3.1-8B，4K-128K 平均）	RULER 平均分	87.79（FlashPrefill V2，BF16）；FP8 变体 86.57	全注意力 88.82；最强稀疏基线 XAttention 87.37	与全注意力差距仅 1.03 分，超过所有稀疏基线约 0.4 分
注意力算子加速 @128K（batch 4，H20）	相对 FlashAttention-2 的加速比	27.19×-27.45×（BF16），47.26×-47.59×（FP8）	FlashPrefill V1 18.67×-19.26×；FlexPrefill 5.18×-5.43×；XAttention 3.42×-3.48×；MInference 2.45×-2.76×	BF16 比前作再快约 1.4×，FP8 翻倍以上；对 FA3/4 对齐稠密基线仍 17.5×/30.5×
LongBench 21 任务平均	平均分	49.31 / 46.96 / 50.73（三模型）	XAttention 48.25 / 45.90 / 49.61	每个模型均列稀疏方法第一，与全注意力差 <0.9 分
SGLang 端到端 TTFT @128K（TP=4）	首 token 延迟与加速比	2.15×-3.40×（BF16），3.66×-4.83×（FP8）；如 Qwen3-30B batch16 从 123.2s 降至 36.2s/25.5s	FA3/4 对齐稠密注意力后端	加速近似 batch 不变，batch 越大收益越稳定
开环服务（Poisson 到达，4K-128K 混合负载）	请求吞吐（req/s）与 P50 TTFT	0.70-0.76 req/s（BF16）、0.88-1.34 req/s（FP8）；P50 TTFT 17-46s，FP8 最低 2.79s	FA3/4 饱和于 0.31-0.37 req/s，P50 TTFT 77-106s	吞吐约 2×-3.6×，低速率下 P50 TTFT 加速 4.5×-27.5×
与生产内核对比 @64K FP8	块稀疏注意力延迟	FlashPrefill V2-BSA，各稀疏度均最优，接近理论界	HPC-Ops BSA（生产推理算子库，仅支持 FP8）	所有稀疏度下快 6-7%，稠密参考两者相当（123.7ms vs 133.4ms）

局限与改进

作者承认的局限：(1) decode 退回稠密注意力（单查询无块级稀疏空间），只加速 prefill；(2) 端到端加速受注意力占 prefill 计算比例约束，4K 时密度约 70%，BF16 与 FA3/4 基本持平（1.01×-1.09×），索引开销抵消收益；(3) chunked prefill 兼容性受损——索引每个 chunk 重跑、强制尾块推高短 chunk 有效密度，8K chunk 时 Qwen3-30B BF16 P50 收益从 2.2× 降到 1.2×，论文建议 chunk ≥8K；(4) 90% 稀疏时校正相对开销升至 18%（BF16）/27%（FP8）；(5) FP8 在 128K 与 BF16 仍有约 0.4-2.3 分差距。我的补充观察：(6) 评估以合成任务与 LongBench 为主，缺少真实生产负载（代码补全、多轮 agent）下注意力模式不可预测时的鲁棒性数据；(7) 仅在 H20/Hopper 验证，未覆盖其他架构；(8) 误差界依赖块内低方差假设，对高方差头缺乏自适应机制；(9) 未系统分析关键块被误剪时的失效行为。

独立分析的弱点

(1) 短序列无收益：4K 密度约 70%，BF16 端到端仅 1.01×-1.09×。改进：长度/密度感知路由，短序列直接走稠密 FA3/4 路径，或用更廉价的轻量探测降低索引成本。(2) chunked prefill 场景索引重算与强制尾块推高有效密度，8K chunk 收益明显侵蚀。改进：跨 chunk 缓存并增量更新 CSR 索引，或按 chunk 长度缩放强制保留比例。(3) FP8 精度损失放大：量化压缩分数边际使无校正时 128K 掉 6.2 分。改进：校正统计用 FP16 池化而计算仍 FP8，或 per-head 自适应 $\alpha$。(4) decode 不受益，长输出场景总时延仍受 decode 主导。改进：与 query-aware KV 选择（Quest 类）共享 CSR 基础设施实现 decode 稀疏。(5) 负载均衡按选中块数均分，tile 间选中数方差大时有长尾。改进：引入代价模型加权分片。(6) 校正在 90% 稀疏时开销 18-27% 且假设块内低方差，对高方差头可自适应增大采样。各弱点均有明确工程抓手，不影响主结论。

未来方向

作者结论聚焦把稀疏注意力推向实用，未列明确 future work，基于成果可延伸：(1) 扩展到 decode 阶段：prefill 生成的 CSR 与池化统计可复用于解码期 query-aware KV 选择，实现全阶段稀疏；(2) 新硬件移植：FA4 针对非对称硬件扩展的协同设计思想可迁移至 Blackwell，并探索 FP4 或更低位宽与稀疏的联合；(3) 训练协同：把均值校正作为训练期正则，让模型对极端稀疏鲁棒，进一步压低密度上限；(4) 自适应稀疏：按层/按头动态调整 $\alpha$ 与密度，替代当前全层统一配置（$B=128$、sink 256、窗口 512、$\alpha=0.1$）；(5) 与 KV cache 量化、MoE、混合注意力架构（如 UniPrefill 的 token 级稀疏化）叠加，做精度-效率联合建模；(6) 系统侧：多机 TP/PP 部署验证、与 chunked prefill 的增量索引协同、扩展到视觉 token 稀疏模式不同的多模态长上下文。

复现评估

论文给出代码仓库 github.com/qhfan/FlashPrefillv2，环境完全公开：CUDA 12.9、PyTorch 2.9.1、CUTLASS 4.3、SGLang 0.5.10，硬件为 NVIDIA H20（端到端用单机 TP=4 四卡）。超参数公开且跨模型统一：选择块 $B=128$、256 sink token、512 局部窗口、$\alpha=0.1$、tile 128×64，对比公平性好。评测覆盖 RULER 与 LongBench 公开基准及三个开源模型（Llama-3.1-8B、Qwen3-4B-2507、Qwen3-30B-A3B-2507）。复现难度中等偏高：内核深度绑定 Hopper 专有特性（TMA、wgmma、warp 专用化、FP8 布局处理），需要 H100/H20 级 GPU 与较强 CUDA/CUTLASS 工程能力；算法部分（均值校正公式、max 阈值选择）可在任意框架独立验证。仓库代码完整度与权重放出情况需以实际仓库为准。

FlashPrefill V2：面向长上下文 LLM 服务的块稀疏 Prefill 注意力 FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving