TLX优化Jagged Flash Attention：Blackwell上逼近SOTA的FA4实践

为什么Jagged Flash Attention值得关注

Flash Attention自诞生以来，已经成为现代大模型注意力计算的事实标准。它通过分块计算与在线softmax，把原本需要显式存储完整注意力矩阵的开销压缩到可控范围，大幅降低了显存占用并提升了吞吐。而在推荐与广告系统这类场景中，序列长度往往参差不齐，标准的定长Flash Attention需要大量padding，造成算力浪费。

Jagged Flash Attention（JFA，锯齿状Flash Attention）正是为解决这一痛点而生。它直接处理变长（jagged）序列，避免为对齐而填充无效token。据原文介绍，JFA是支撑Meta生成式广告模型（Generative Ads Model，GEM）背后的核心注意力内核，可见其在大规模工业级推荐/广告系统中的实际价值。

Flash Attention的核心创新在于"分块（tiling）+ 在线softmax（online softmax）"范式。传统注意力计算需要先在HBM（高带宽显存）中完整写入形如 N×N 的注意力分数矩阵，再做softmax归一化，显存复杂度为 O(N²)。Flash Attention将Q/K/V矩阵切成小块逐块装入SRAM（片上缓存），利用数值稳定的在线softmax累积中间统计量（行最大值与指数和），无需落盘完整矩阵，将显存复杂度压缩到 O(N)，且大幅减少HBM读写次数——后者正是GPU注意力计算的主要瓶颈。目前FA1、FA2、FA3到FA4的演进主线是不断提升对硬件异步特性（Tensor Core流水线、异步内存拷贝）的利用率，使实际吞吐逼近理论峰值FLOPS。

在Blackwell(B200)上重写内核的意义

NVIDIA Blackwell架构（B200）带来了新一代张量核心与更高的内存带宽，但要充分释放其算力，往往需要针对新架构重新设计内核的调度与数据流。原文的工作重点，就是把JFA移植并优化到Blackwell平台上，目标是向当前业界最优（SOTA）的FA4（Flash Attention 4）水平靠拢。

这里有两层含义值得拆解：

架构迁移不是简单的编译切换

从上一代GPU迁移到Blackwell，并不是重新编译一遍就能拿到性能。新架构的内存层级、异步拷贝机制、张量核心指令集都在变化，内核的分块策略、流水线安排、寄存器分配都需要重新调优。对于JFA这种本身就要处理变长序列、分支逻辑更复杂的内核，迁移难度更高。

对标FA4意味着对标最高标准

Flash Attention系列每一代都在逼近硬件理论峰值。把一个处理锯齿序列的特化内核做到接近FA4的效率，意味着既要保留变长处理带来的算力节省，又不能因为额外的控制开销拖累峰值吞吐——这是工程上相当有挑战的平衡。

Flash Attention 4（FA4）是专为NVIDIA Hopper/Blackwell架构设计的最新版本，其核心改进包括：充分利用Hopper引入的Tensor Memory Accelerator（TMA）进行异步数据搬运、将计算与数据加载完全流水线化（producer-consumer warp分离）、以及针对FP8低精度路径的专项优化。在H100上，FA4的前向传播可达约接近理论峰值的75%以上的MFU（Model FLOP Utilization）。将一个额外引入变长分支逻辑的JFA内核做到接近FA4的效率，意味着变长带来的控制流开销需要被完全隐藏在异步流水线中，这正是Blackwell重写工作的核心挑战所在。

TLX在内核优化中扮演的角色

原文标题点名使用TLX来构建与优化这一内核。TLX作为一种面向底层GPU内核开发的工具/语言层，能够让开发者以更高的抽象描述张量计算，同时保留对调度和内存访问的精细控制。

相比直接手写CUDA，这类工具的价值在于：

• 降低针对新架构（如Blackwell）重写内核的工程成本
• 让变长序列这类不规则计算模式更容易表达
• 在保持接近手工优化性能的同时提升可维护性

对于需要频繁跟随硬件迭代而重写内核的团队来说，这种可移植性与表达力往往决定了能否快速吃到新硬件的红利。

TLX（Tensor Language X，或类似定位的DSL）属于"内核编译器/领域特定语言"这一技术赛道，与Triton（OpenAI）、Hidet（微软）等工具同属一类思路：在CUDA之上提供更高层的张量操作抽象，由编译器负责将其映射到具体架构的指令（如Blackwell的新一代异步warpgroup MMA指令），开发者无需手动处理寄存器分配、bank conflict规避、流水线插槽等极度依赖硬件细节的工作。这类工具与纯手写CUDA的差距正在随编译器成熟度提升而收窄，但对于变长/稀疏等控制流复杂的场景，高层抽象对工程效率的提升尤为明显，因为手工处理这类不规则访存模式的正确性与性能调优成本极高。

工业界注意力优化的大趋势

从这篇工作可以读出几个行业信号。其一，注意力内核的竞争已经从通用LLM扩展到推荐、广告等结构化稀疏/变长场景，变长优化正成为重要战场。其二，硬件代际更替（Blackwell）正在推动内核层的全面重写，谁能更快把特化内核迁移到新架构，谁就能在成本与延迟上占得先机。其三，像TLX这样的内核开发工具，正在成为连接算法创新与硬件能力的关键基础设施。

对于从事大规模推理与训练系统的工程团队，这类工作的启示很直接：不要把注意力内核当成黑盒，针对业务的序列分布特征做特化，往往能换来显著的成本节约；同时，提前投资于可移植的内核开发工具链，能让团队在每一次硬件升级中更快兑现性能。

小结

这篇工作把Meta GEM背后的Jagged Flash Attention内核搬到NVIDIA Blackwell平台，借助TLX进行优化，目标是逼近SOTA级别的FA4性能。它既展示了变长注意力在工业场景的现实需求，也反映出内核优化正在随硬件迭代而持续演进。由于原始素材为预告式摘要，关于具体的性能数据、实现细节与对比基准仍有待完整技术博客披露。