📖 AI 总结

Meta 的生成式广告模型 GEM 需要在长度不一的用户序列上运行注意力计算，若将序列填充至固定长度会浪费多达一半的算力，因此采用紧凑打包的锯齿状表示，并由此产生了 Jagged Flash Attention 这一核心内核。该内核是 GEM 中最慢的环节，此前要在 NVIDIA Blackwell（B200）上达到峰值性能，通常只能依赖手写 CuteDSL 或 CUDA，开发周期长且难以扩展。本文介绍用 TLX（Triton 底层扩展）重写该内核的成果：代码约 3200 行，比 FlashAttention-4 约 1 万行的 CuteDSL 实现少约三倍；在 GEM 关注的非规则形状上，前向性能提升约 13%，反向提升约 50%。由于仍处于 Triton 高层编程模型，该内核更易被建模工程师阅读和扩展，为在 Blackwell 上实现接近 SOTA 的注意力内核提供了兼顾效率与性能的路径。