企业官网建设流程全解析

hpc-ops是腾讯混元 AI Infra 团队开源的高性能 LLM 推理算子库（High Performance LLM Inference Operator Library）。其中的块稀疏 FP8 Attention（Block-Sparse FP8 Attention）专为长上下文 Prefill 设计：通过预计算的块级掩码直接跳过无关 KV 块，并在 FP8 精度下保持数值质量，在不同稀疏度与序列长度下相比 MIT-BSA BF16、FlashPrefill-BSA BF16、HPC-Dense FP8 和 FA3-Dense FP8 等基线，最高可取得3.16 倍加速。

🐢 长上下文 Prefill 的瓶颈：大部分 KV 计算是"白做的"

长文档推理时，Prefill 阶段需要让每个新 Query 关注 KV Cache 中的历史 token。但现实是：大部分 KV 内容与当前 Query 相关性极低——例如检索式注意力、滑动窗口模式或长文档中的跨章节噪声注意力。

这带来一个关键事实：Prefill Attention 在长序列下通常是访存带宽受限的，而稀疏掩码（block-sparse mask）可以让整个 KV tile 都不参与计算。块稀疏 Attention（BSA）的价值就在于：不是"算完再屏蔽"，而是"根本不读、根本不算"被掩码掉的 KV tile。

hpc-ops 的块稀疏 FP8 Attention 同时解决了三个难题：

🧩 块稀疏调度：block_mask 直接跳过 KV Tile

核心 API 是 hpc/attention.py 中的attention_with_kvcache_blocksparse_prefill_fp8。它接收一个block_mask张量：

• Shape：[num_batch, num_head_q, max_tile_m, Kb]，其中每个 Q-tile 为 128 个 token，每个 KV-tile（kTileN）为 128 个 token，Kb = ceil(max_kv_len / 128)
• Dtype：uint8，True = 计算，False = 跳过
• 不传 mask 时，走同一 kernel 家族的kHasMask=False稠密路径，因此该接口同时是稠密/稀疏的统一入口

kernel 内部（src/attention/prefill/sm90/kernels.cuh）用 128×128 的 tile 组织 Q 与 KV，每个 CTA 只驻留一个 SM（384 线程 = 1 个生产者 warpgroup + 2 个消费者 warpgroup，典型 warp specialization 结构）。

活跃 tile 列表：一次压缩，全流水线受益

如果每个消费者都要逐 tile 查 mask，稀疏判断本身会成为串行瓶颈。hpc-ops 的做法是：

1. 生产者 warpgroup 的 warp 0用__ballot_sync做一次 warp 并行的流压缩（stream compaction），把该 Q-tile 下所有"活跃" KV-tile 的索引顺序写入 SMEM 的shm_active_tiles数组；
2. 生产者随后只对活跃 tile 发起 TMA 加载 K/V（跳过的 tile 完全不触碰 HBM）；
3. 消费者按同一份列表逐 tile 执行 QK^T → online softmax → PV 累加。

这意味着稀疏度 50% 时，KV 的全局内存读取量近似减半，而调度开销被压缩进一次 warp 级 ballot，几乎可以忽略。

⚡ FP8 Tensor Core 计算：固定 P-scale 量化技巧

两个 GEMM（QK^T 与 PV）都运行在 SM90 FP8（e4m3）Tensor Core 上（SM90_64x64x32与SM90_64x128x32GMMA 原子，见 warp_spec_with_kvcache_blocksparse_fp8_dim128.cu）。

softmax 得到的 P 矩阵是动态值，直接转 FP8 会丢失精度。hpc-ops 采用固定 P-scale = 256的技巧：

• P 累加器先乘以 256 再量化为 FP8，把动态范围整体平移进 e4m3 的安全区间；
• 末尾 epilogue 用倒数kFp8PrefillPScaleInv一次性还原，与vscale合并为一个乘加；
• 因果边界外（icol > irow）的注意力分数置为-inf，online softmax 自动处理。

接口支持两种量化方案（QuantType）：

• QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR：K/V 全张量缩放，实现最简；
• QPERTOKEN_PERHEAD_KPERTOKEN_PERHEAD_VPERHEAD：K 按 token 组 + head + 维度组细粒度缩放，V 按 head 缩放，精度更稳。

🚀 3.16 倍加速是怎么来的

README 性能表（README.md）中块稀疏 FP8 Attention 一行：

加速由两部分叠加：FP8 带来的计算/访存减半（相对 BF16 基线），加上块稀疏跳 tile 带来的带宽节省（相对稠密 FP8 基线）。序列越长、稀疏度越高，跳过的 KV 越多，3.16x 正是在长序列 + 高稀疏比下测得的峰值。

正确性由测试保证：tests/test_attention_blocksparse_qpertoken_perhead_kvpertensor_fp8.py 与 tests/test_attention_blocksparse_qkpertoken_perhead_vperhead_fp8.py 分别在skip_ratio = 0.0 / 0.5下与 naive 参考实现逐位对拍（atol=0.1），并覆盖 NHD/HND 两种 KV cache 布局与 32/64 两种 page block 大小。

🛠️ 快速上手：一行调用块稀疏 FP8 Attention

import hpc import torch out = hpc.attention_with_kvcache_blocksparse_prefill_fp8( q, kcache, vcache, qscale, kscale, vscale, cu_seqlens_q, block_ids, seqlens_kvcache, max_seqlens_q, quant_type=hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR, block_mask=block_mask_u8, # uint8，True=计算 )

使用注意（源自 src/attention/entry.cc 的参数校验）：

• 仅支持dim_qk = dim_v = 128，K/V cache pageblock_size为 32 或 64；
• 建议每个 Q-tile 因果范围内的对角 tile 置 True，否则全-infsoftmax 会出 NaN（kernel 本身保证任意 mask 不死锁）；
• KV cache 为 paged 格式（page table 由block_ids提供），未用槽位需零填充；
• 当前 sm90（H20/H800 等）架构支持，由HPC_ARCH_DISPATCH精确分发。

📂 源码导航

总结

hpc-ops 的块稀疏 FP8 Attention 用一套"活跃 tile 列表 + TMA 只加载 + 固定 P-scale 量化"的组合拳，把长上下文 Prefill 中最浪费的 KV 访存直接砍掉，同时在 FP8 精度下保住数值质量——这就是相对 BF16 块稀疏基线最高3.16 倍加速背后的完整原理。对于长文档、RAG、代码库理解这类 Prefill 密集场景，它是目前开源方案中非常值得关注的稀疏 Attention 实现。