行业动态
dev.to ·
@stmanst
SFPU 内部:一个 40 年前的舍入技巧为何在现代 AI 加速器上失效
作者在 Tenstorrent 的 ML 框架 tt-metal 中调试 SFPU(负责 exp、log、softplus 等超越函数的硬件数学单元)时,发现 exp(x)*exp(-x) 这类数学上恒为 1 的计算链会因中间结果超出 fp32 范围而溢出,修复方式是重排计算图、用代数恒等式避免存储超范围中间值。文章还指出 SFPU 使用 IEEE 754 默认的 round-to-nearest-even,在 fp16/bf16 中间计算转回 fp32 时与 PyTorch 参考值存在可见差异,并建议用包含 ±0.0、±inf、次正规数等边界输入对照参考实现测试。
AI推荐理由 作者以在 tt-metal(Tenstorrent ML 框架)上的实际调试经历,描述 SFPU 数学单元中 exp(x)*exp(-x) 链式运算的中间结果溢出问题,以及通过重排计算图规避溢出、并指出 round-to-nearest-even 在 fp16/bf16 转回 fp32 时与 PyTorch 参考值产生偏差。属于具体硬件数值实现细节的一手经验,对做加速器算子与数值对齐的开发者有参考价值;但内容为个人博客叙述,未给出可复现的测试数据或补丁链接,结论范围有限。
72
正文
The Hidden World of SFPU Rounding
In my work on tt-metal (Tenstorrent's ML framework), I encountered a subtle but critical bug in the SFPU (SFPU = Tensor Processing Unit math unit) that caused intermediate overflow in floating-point computation chains. Here's what I found.
What is the SFPU?
The SFPU (StochaSTic Processing Unit, or more likely the hardware math unit) handles transcendental functions like exp, log, and softplus on Tenstorrent chips. These functions use polynomial or rational approximations because the hardware does not implement them directly.
The Overflow Problem
Consider a chain of operations: exp(x) * exp(-x). Mathematically this equals 1 for all x. But in SFPU computation:
exp(x) is computed and stored as an intermediate result
exp(-x) is computed
The multiplication exp(x) * exp(-x) overflows if the intermediate result exceeds the fp32 range
The fix was to reorder computation to avoid storing intermediate results that exceed the representable range. Instead of computing and storing both exp values, we restructure the computation graph to use algebraic identities that eliminate the overflow.
The Rounding Trick
The SFPU uses a rounding mode called round-to-nearest-even (IEEE 754 default). This is correct, but when converting back to fp32 after intermediate fp16/bf16 computation, the rounding can cause small but visible differences from PyTorch reference values. The trick: add 0.5 before truncation for round-to-nearest, rather than using the hardware rounding mode.
This is a 40+ year old optimization from IEEE 754, but it interacts subtly with the SFPU's internal precision.
The Real Lesson
Hardware-accelerated ML frameworks are full of these subtle numerical bugs. The fix required understanding both the mathematical properties of the functions and the hardware-level implementation details. Always test against the reference PyTorch implementation with a wide range of inputs — including edge cases like ±0.0, ±inf, and subnormal numbers.
Follow my bug bounty journey: @truongsontung
Related Posts
Inside SFPU Overflow Bugs
Building an Autonomous Bug Bounty Hunter