PyTorch 和 Meta 的广告排序团队本周丢出了一个 Blackwell 专用的注意力内核——TLX Block Attention——更大的故事是它下面的那一层。**TLX(Triton Language Extensions)**是 Triton 的生产力与 Blackwell 上原始 CUTLASS 级控制之间的 DSL 桥梁,把新的 tcgen05 异步张量核、TMA 描述符和 TMEM(每 SM 256KB 张量内存)暴露为 Triton 原语,如 `tlx.async_dot`、`tlx.async_descriptor_load`、`tlx.local_trans`,加上用于生产者-消费者 warp 管线的 mBarrier 同步。Repo:github.com/triton-lang/triton-ext。这是 2026 年大多数编写 Blackwell 内核的 builder 在它稳定后将居住的层。
内核本身针对固定块稀疏自注意力——64-token 块,块对角模式,编译时已知。这个形状专门用于 Meta 的广告排序和推荐模型,不是 LLM 注意力。因为模式在编译时已知,内核消除了 Flash Attention 的多 tile 迭代循环、在线 softmax 修正因子、logsumexp HBM 往返,以及单独的 Di 预处理——每个 Q tile 恰好关注一个 K/V tile,单个 GEMM,无需修正。正向通过每 CTA 使用 15 个 warps 在专门管线中(1 load / 1 QK-MMA / 4 softmax / 1 PV-MMA / 8 epilogue);反向跨 7 个阶段使用 20 个 warps。正向 TMEM 三缓冲(~169KB / 256KB),反向双缓冲(~162KB / 256KB)。B200、BF16、稀疏度=70%——正向 0.98ms vs Flash Attention v2 的 1.81ms(1.85 倍),反向 2.36ms vs 5.89ms(2.50 倍),总计 2.31 倍。数值精度在 max dQ diff 上比 FA v2 好 53%。
融合 rotary 反向是第二个亮点和可推广的模式。独立 attention 反向 1.56ms 加 rotary 反向 4.88ms = 6.44ms 未融合;融合成单一内核,在 TMEM/寄存器中保持 dV 为 FP32,就地应用 rotary 共轭,然后做一次 BF16 全局 store = 1.82ms。**快 3.54 倍。**这条经验在广告工作负载之外是可移植的:当你在寄存器/TMEM 中有 FP32 中间值时,以 FP32 做你的 epilogue 数学并 store BF16 一次,消除了原本主导的通过全局内存的往返。即使没有 TLX 或 Blackwell,builder 也可以将这种洞见应用到其他融合操作内核。
周一上午:这个内核按发货状态对你有用,如果你在 B200/B300 GPU 上交付带块对角注意力的广告排序、推荐或特征交互模型——克隆 facebookresearch/ads_model_kernel_library 并 benchmark。如果你是 LLM builder,该内核不适用(因果、滑动窗口和任意稀疏模式被明确排除),但 TLX DSL 本身是值得关注的部分——这就是 Blackwell 感知的 Triton 内核将如何被编写,大多数架构原语(warp 专门化、TMA 描述符、TMEM 累加器)适用于你的栈所需的任何注意力形状。诚实的限制:仅 Blackwell(sm_100+),无 Ampere/Hopper 后备,head_dim 硬编码为 64 或 128,块大小 64 固定,博客中未声明许可证(检查仓库)。对于使用这些技术的 LLM 形状注意力,Flash Attention 3 的 Blackwell 移植版及其后继版本将是下季度的观察项。