独立研究者修复 Pallas TPU 已知 bug：位级正确且快 1.67 倍
Francis_YAO_ · x · 2026-10-11
独立研究者 Muhan Zhong 发布预印本，证明 Pallas TPU 官方 top-k 实现中一个「因修复成本高而故意保留」的 bug 其实便宜又值得修。
- 问题：在 TPU 上 Pallas 内核里调用 jax.lax.topk 时，输入含 -inf 会返回重复索引；源码注释称修复太贵而有意保留。实测其错误不止重复索引：197592 个含特殊值的输入中有 70706 个结果错误。
- 修复成本：以 IEEE 754 totalOrder 定义正确性，位级正确的公式在 f32[8,128] 取 top-8 仅多花 4.7% 周期——所谓「太贵」并不成立。
- 进一步提速：将运行时间建模为跨 lane 单元 XLU 的往返延迟、发射间隔与逐元素向量指令数，按形状选择公式；仅用现有 API，在 25 个形状中的 20 个上快于官方 Pallas（几何平均提速 1.67×），且全面快于原生 XLA。
- 对剩余形状提出 fold-and-rank 与全转置后 loser-tree 合并两种算法（需用 Pallas 目前无法触达的指令），可将领先形状数提高到 24。
「Infra」频道最新
- Let's Encrypt 宣布 2027 年 2 月起默认签发 64 天证书 — jedisct1 · 2026-10-11
- Anthropic 工程师自曝单日消耗 10 亿 token,连跑一周 — AaronBergman18 · 2026-10-11
- 访谈 NVIDIA 推理专家：GPU 退役时机与 AI 基础设施经济账 — kimmonismus · 2026-10-11
- AI agent 调优 Triton kernel，翻转 grid 顺序换来 1.29x 推理提速 — zmkzmkz · 2026-10-11
- 囤了 12 万美元 GPU 的 ML 工程师:只因好玩,不劝人跟风 — TheZachMueller · 2026-10-11
- 迁移 AI 工作负载常意味着 70% 以上重写,成本远超预期 — DavidLinthicum · 2026-10-11