PyTorch 用 Triton 替换 CUDA，推荐系统 kernel 前向提速 1.28 倍

PyTorch · x · 2026-10-07

PyTorch 官方发布工程深度解析：Meta 团队将推荐系统中核心的 Table Batched Embedding（TBE）kernel 从 CUDA 迁移到 Triton（FBTriton），取得显著收益。

关键结果

• 前向传播最高提速 1.28x，反向传播最高 2x，开发者效率大幅提升
• TBE 在单次 GPU launch 内完成多张 embedding 表的查找与池化，降低 launch 开销、提升内存效率

技术要点

• 前向实现分两条路径：通用 gather 路径 + 基于小表直方图的快速路径
• 通过调优每个 program 处理的 bag 数（大负载用 2 bags/program）、四次独立行加载的 gather 宽度等手段压榨性能
• 支持 int32 索引/偏移（范围 < 2^31 时），索引存储减半

文章还指出了后续优化方向。对做推荐系统或多卡 embedding 服务的工程团队有直接参考价值。

原文链接 →

「Infra」频道最新

• AI 需求压垮内存供应，三星 Q3 利润预计同比暴涨 770% — Polymarket · 2026-10-07
• 一人项目 openTPU:用 AI Agent 设计出开源加速器,FPGA 上跑通多个 LLM — ai · 2026-10-07
• 免费线上会议 All Day AI 定档 10 月 22 日,聚焦小模型与本地推理 — FikoFox · 2026-10-07
• 4000 美元预算跑大模型：R9700、Strix Halo 还是 Arc B60？ — BinaryGrind · 2026-10-07
• 巴黎蒙帕纳斯大楼拟装 12 万块 B300,540 EFLOPS 打造欧洲算力灯塔 — IgorCarron · 2026-10-07
• 实测：27B 模型塞进 12GB 显存+8GB 内存，仍有约 18 tok/s — bodhi371 · 2026-10-07

更多「Infra」频道 AI 资讯 →