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
Top comments (0)