测试表达能力,不只测试编译器内部
mapping_gemm.py 的源码首先声明行、列、归约域,再把输入、逻辑 buffer、storage、stage、fill 和 contraction 连起来。它没有把 launch geometry 当成用户程序。
相应的 N64 QK 源码测试还检查,用户程序的序列化表示中不出现 ExecutionPlan、KernelSpec、gpu.launch 或 WGMMA 等后端输入。同时,解析结果可以选择具体的 SM90 WGMMA Atom。作者语义与目标实现由此分别观察。
Attention 是一条组合链
register_attention.py 为 Q、K、V、scores、maximum、exponentials、sums 和 probabilities 分别建立逻辑值。源码说明其目标是让概率 tile 留在寄存器表示中,再进入 register-A/shared-B 的 PV contraction。
这里重要的不是只出现了两次矩阵乘。中间值的 dtype、状态、存储交接和最终发布都属于同一条数据流。完成了一次 QK 并不能证明整个 attention 语义已闭合。
从一个 tile 走向在线递推
FlagKernel README 把 FA3 路径分成逐级证明:先 launch 与 copy,再单 tile GEMM、row max/sum、指数归一化、tile softmax,最后连接 QK→softmax→PV。多个 K tile 还需要继续维护在线 m/l/O 状态。
进入流水线阶段后,双缓冲、causal/tail 条件和传输时序又带来新的要求。证明两个 block 写互不重叠的输出,与证明它们通过跨 block 通信合作计算一个输出,也不是同一件事。
四种证据不能互相替代
| 观察 | 能够说明什么 |
|---|---|
| 源码/IR 检查 | 关系被表达,或某类非法程序被拒绝。 |
| 生成代码检查 | 目标指令、资源或时序机制确实生成。 |
| 设备数值测试 | 限定输入和容差下与独立参考一致。 |
| 匹配条件的 benchmark | 指定设备、精度、形状及计时边界下的性能。 |
README 保留了多个历史 H20 结果,但本篇不把它们合并成一个“整体比其他框架更快”的结论。普通 launch 与 CUDA Graph、带尾部 mask 与去掉 mask 的任务,都不能互换比较。
产品价值来自反复使用同一条入口
每一个新增真实算子都在追问:作者能否只用公开 API 写出它?失败时能否回到对应源码?状态和内存关系能否被检查?生成的实现能否按相同数值语义运行?
这让 FlagKernel 成为语言改进的反馈来源。当前仓库中有多种 attention 和 mapping 程序,并不意味着任意生产形状、mask、设备或组合均已获得资格。明确这一点,反而更容易识别下一项最值得完成的语言能力。
从算子表达走到性能结果
FlagKernel 已有与 FlashInfer FA3 的直接性能对比。P0be 的八种 dense attention 规模全部领先,延迟几何平均降低 10.33%;P0z 的一个 batch/GQA/tail 场景在八张 H20 配对测试中,延迟降低约 43.5%。这些结果让公开 DSL 的表达能力与编译器生成的执行方案形成了完整案例。
源码依据
基于本地开发快照撰写。下列为源码中的相对路径与仓库提交号,不是公开下载链接;本文没有重新执行 GPU 测试,也不代表冻结发布包的全部能力。
FlagKernel/README.md6c31111666FlagKernel/python/flagkernel/mapping_gemm.py6c31111666FlagKernel/python/flagkernel/register_attention.py6c31111666FlagKernel/test/test_n64_qk_source.py6c31111666