← FLAGLANG / 研发专栏

FlagKernel:让真实算子检验语言

从公开 DSL 编写 GEMM、softmax 和 attention,连接语言表达与 GPU 实现。

测试表达能力,不只测试编译器内部

mapping_gemm.py 的源码首先声明行、列、归约域,再把输入、逻辑 buffer、storage、stage、fill 和 contraction 连起来。它没有把 launch geometry 当成用户程序。

相应的 N64 QK 源码测试还检查,用户程序的序列化表示中不出现 ExecutionPlan、KernelSpec、gpu.launch 或 WGMMA 等后端输入。同时,解析结果可以选择具体的 SM90 WGMMA Atom。作者语义与目标实现由此分别观察。

Attention 是一条组合链

Q × K→row max / exp / sum→normalize→P × V

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.md6c31111666
  • FlagKernel/python/flagkernel/mapping_gemm.py6c31111666
  • FlagKernel/python/flagkernel/register_attention.py6c31111666
  • FlagKernel/test/test_n64_qk_source.py6c31111666
LET’S BUILD / CLASS OF 2027

一起把系统做深。

AI 编译器 · GPU Kernel · 大模型推理系统

邮件
联系 ↗