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 固定,部落格中未聲明授權(檢查 repo)。對於使用這些技術的 LLM 形狀注意力,Flash Attention 3 的 Blackwell 移植版及其後繼版本將是下季的觀察項。