Non-uniform breakpoint spacing for exp2 piecewise-linear approximation

Investigate non-uniformly spaced breakpoint placement for the piecewise linear approximation of the exp2 function used in FSA’s SystolicAttention to implement the element-wise exponential within FlashAttention, and determine whether such non-uniform spacing reduces mean relative error compared to the current uniformly spaced eight-segment scheme.

Background

FSA implements the element-wise exp operation in FlashAttention as exp2 and approximates it using a piecewise linear interpolation executed within the systolic array. By splitting x into integer and fractional parts (x = x_i + x_f) and exploiting that x_f ∈ (-1, 0], the design uses uniform segments and streams slope/intercept coefficients to reuse MAC units for interpolation.

The authors evaluate the approximation over negative normal fp16 values and select an 8-segment uniform scheme, reporting MAE ≈ 1.4e-4 and MRE ≈ 2.728e-2. They explicitly note that exploring non-uniformly spaced breakpoints to further reduce MRE is left for future work, making the evaluation and design of such schemes an unresolved question in their framework.