PyTorch 用 TLX 写 JFA 内核,前向快 FA4 约 13%、反向快 50%
PyTorch · x · 2026-10-02
PyTorch 团队发布博客，介绍为 Meta 生成式广告模型 GEM 打造的注意力内核 Jagged Flash Attention（JFA）在 NVIDIA Blackwell B200 上的实现与优化。
- 实现方式：基于 TLX（Triton Low-level Extensions）构建，在 Triton 的高层 tile 编程模型之上增加了显式的硬件级控制。
- 代码量优势：约 3.2K 行简洁的 Triton 级代码，比 SOTA 的 FlashAttention-4（FA4，约 10K 行 CuteDSL）少约 3 倍。
- 性能表现：在 GEM 关注的 jagged 形状上超过 FA4——前向传播快约 13%，反向传播快约 50%。
文章展示了用高层语言写接近手写性能内核的可行性，对关注推理/训练内核优化的工程师有直接参考价值。
「Infra」频道最新
- 投资人吐槽:投了钱还得求 GPU 厂给被投公司明年配额 — mattturck · 2026-10-02
- Lambda 完成超 10 亿美元 GPU 债务融资，固定利率 6.78% — TheZachMueller · 2026-10-02
- 双 DGX Spark 集群对战 M5 Ultra Mac Studio 实测视频出炉 — AIFlow_ML · 2026-10-02
- 9070 XT 上 SageAttention 提速近 2 倍：作者手写 HIP kernel 并放出 wheel — Familiar_Worry332 · 2026-10-02
- 网友称自托管 GLM 模型在双 RTX 6000 Pro 上跑出 175 tok/s、98% 草稿接受率 — HankYeomans · 2026-10-02
- Cloudflare 开放账户 API Token 自助创建,CLI 一条命令搞定 — irvinebroque · 2026-10-02