fix act_quant_kernel (#968)

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao 2025-08-27 16:23:30 +08:00 committed by GitHub
parent b15f0dbbbe
commit 82f6008c8c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -23,8 +23,7 @@ def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, scale_fmt: t
pid = tl.program_id(axis=0)
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
x = tl.load(x_ptr + offs).to(tl.float32)
amax = tl.max(tl.abs(x))
amax = tl.min(amax, 1e-4)
amax = tl.max(tl.abs(x), 1e-4)
s = amax / 448.
if scale_fmt == "ue8m0":
exp = tl.math.ceil(tl.math.log2(s))