Skip to content

Commit d5eb18f

Browse files
authored
fix(kernel): correct fp4 cache scale rounding (#720)
* fix(kernel): preserve fp4 scale rounding at subnormal ties * fix(kernel): preserve native e4m3 scale calculation
1 parent 3886cee commit d5eb18f

1 file changed

Lines changed: 6 additions & 1 deletion

File tree

  • python/freetoken/kernel/triton/dsv41

‎python/freetoken/kernel/triton/dsv41/pack.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,12 @@ def _pack_rows_kernel(
8787
amax = tl.maximum(tl.max(tl.abs(xe), axis=1), tl.max(tl.abs(xo), axis=1))
8888
if FMT == 2:
8989
amax = tl.maximum(amax, 6.0 * (2.0 ** -9))
90-
s = _e4m3_round(amax * (1.0 / 6.0))
90+
if e4m3_native_cx():
91+
s = amax * (1.0 / 6.0)
92+
else:
93+
# Round before the emulated magic-add; fusing it changes subnormal ties.
94+
s = tl.math.div_rn(amax, 6.0)
95+
s = _e4m3_round(s)
9196
tl.store(dst + D // 2 + offs_g, _e4m3_codes(s))
9297
else:
9398
amax = tl.maximum(amax, 6.0 * (2.0 ** -126))

0 commit comments

Comments
 (0)