PyTorch et l'équipe ads-ranking de Meta ont sorti cette semaine un kernel d'attention spécifique Blackwell — TLX Block Attention — et la plus grosse story c'est la couche en-dessous. **TLX (Triton Language Extensions)** c'est le pont DSL entre la productivité de Triton et le contrôle niveau CUTLASS raw sur Blackwell, exposant les nouveaux tensor cores async tcgen05, les descripteurs TMA et la TMEM (256KB-par-SM Tensor Memory) comme primitifs Triton comme `tlx.async_dot`, `tlx.async_descriptor_load`, `tlx.local_trans`, plus la synchronisation mBarrier pour les pipelines warp producer-consumer. Repo : github.com/triton-lang/triton-ext. C'est la couche dans laquelle la plupart des builders qui écrivent des kernels Blackwell en 2026 vont vivre une fois qu'elle se stabilise.
Le kernel lui-même target l'attention self-attention sparse à blocs fixes — blocs de 64 tokens, pattern block-diagonal, connu à la compilation. Cette forme c'est spécifiquement pour les modèles ads-ranking et recommendation de Meta, pas l'attention LLM. Parce que le pattern est connu à compile time, le kernel élimine la loop d'itération multi-tile de Flash Attention, les facteurs de correction d'online-softmax, le round-trip logsumexp HBM, et le preprocessing Di séparé — chaque tile Q attend à exactement un tile K/V, un seul GEMM, pas de correction nécessaire. Le forward pass utilise 15 warps par CTA dans un pipeline spécialisé (1 load / 1 QK-MMA / 4 softmax / 1 PV-MMA / 8 epilogue) ; le backward utilise 20 warps à travers 7 stages. TMEM triple-buffered en forward (~169KB / 256KB), double-buffered en backward (~162KB / 256KB). B200, BF16, sparsity=70% — forward 0,98ms vs 1,81ms de Flash Attention v2 (1,85×), backward 2,36ms vs 5,89ms (2,50×), total 2,31×. La précision numérique bat FA v2 de 53% sur max dQ diff.
Le rotary backward fusé c'est le second highlight et le pattern généralisable. Attention backward standalone 1,56ms plus rotary backward 4,88ms = 6,44ms non-fusé ; fusé en un seul kernel qui garde dV en FP32 dans TMEM/registres, applique le conjugué rotary in-place, puis fait un seul store global BF16 = 1,82ms. **3,54× plus rapide.** La leçon est portable au-delà des workloads ads : quand t'as des valeurs intermédiaires FP32 dans registres/TMEM, faire ta math d'epilogue en FP32 et store BF16 une fois élimine les round-trips à travers la mémoire globale qui autrement dominent. C'est le genre d'insight que les builders peuvent appliquer à d'autres kernels d'ops fusés même sans TLX ou Blackwell.
Lundi matin : ce kernel tel que shipped est utile pour toi si tu shippes des modèles ad-ranking, recsys ou feature-interaction avec de l'attention block-diagonal sur des GPUs B200/B300 — clone facebookresearch/ads_model_kernel_library et benchmark. Si t'es un builder LLM, le kernel ne s'applique pas (causal, sliding-window, et patterns sparses arbitraires sont explicitement exclus), mais le DSL TLX lui-même c'est la partie à watcher — c'est comment les kernels Triton Blackwell-aware vont être écrits, et la plupart des primitifs architecturaux (warp specialization, descripteurs TMA, accumulateurs TMEM) se généralisent à n'importe quelle shape d'attention que ta stack a besoin. Limits honnêtes : Blackwell-seulement (sm_100+), pas de fallback Ampere/Hopper, head_dim hardcoded à 64 ou 128, block size 64 fixé, licence pas mentionnée dans le blog (check le repo). Pour l'attention LLM-shaped avec ces techniques, le port Blackwell de Flash Attention 3 et ses successeurs sera le watch item du prochain trimestre.