[Bugfix][ROCm] Use a zero default for masked scales in the MXFP8 GEMM (#59454)

Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com>
This commit is contained in:
Mehmet Cagri
2026-10-01 11:32:25 -05:00
committed by GitHub
parent 9aaef4296f
commit 08e03df940
@@ -145,16 +145,16 @@ def _block32_tiled_kernel(
for kk in range(0, K_PER_SPLIT, BLOCK_K):
if EVEN_K:
x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=127)
xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=127)
ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
else:
k_ok = (offs_k + kk) < K
s_ok = (offs_sk + kk // 32) < K // 32
x = tl.load(x_ptrs, mask=m_mask[:, None] & k_ok[None, :], other=0.0)
xs = tl.load(xs_ptrs, mask=m_mask[:, None] & s_ok[None, :], other=127)
xs = tl.load(xs_ptrs, mask=m_mask[:, None] & s_ok[None, :], other=0)
w = tl.load(w_ptrs, mask=n_mask[:, None] & k_ok[None, :], other=0.0)
ws = tl.load(ws_ptrs, mask=n_mask[:, None] & s_ok[None, :], other=127)
ws = tl.load(ws_ptrs, mask=n_mask[:, None] & s_ok[None, :], other=0)
acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
x_ptrs += BLOCK_K
w_ptrs += BLOCK_K
@@ -240,9 +240,9 @@ def _block32_packed_kernel(
ws_mask = ws_mask & (wg < K // 32)
x = tl.load(x_ptr + rows[:, None] * K + xk, mask=x_mask, other=0.0)
w = tl.load(w_ptr + cols[:, None] * K + wk, mask=w_mask, other=0.0)
xs = tl.load(xs_ptr + rows[:, None] * (K // 32) + xg, mask=xs_mask, other=127)
xs = tl.load(xs_ptr + rows[:, None] * (K // 32) + xg, mask=xs_mask, other=0)
ws = tl.load(
ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=127
ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=0
)
acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
# Keep only products of matching K panels: the block diagonal.