[Bugfix][ROCm] Use a zero default for masked scales in the MXFP8 GEMM (#59454)
Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com>
Mehmet Cagri committed
Oct 1, 2026 at 11:32 UTC
08e03df94028d5e7cff96bf569459d8d2b926d9e
1 file changed
+6
-6
vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
+6
-6
@@ -145,16 +145,16 @@ def _block32_tiled_kernel(
145
for kk in range(0, K_PER_SPLIT, BLOCK_K):
146
if EVEN_K:
147
x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
148
- xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=127)
148
+ xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
149
w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
150
- ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=127)
150
+ ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
151
else:
152
k_ok = (offs_k + kk) < K
153
s_ok = (offs_sk + kk // 32) < K // 32
154
x = tl.load(x_ptrs, mask=m_mask[:, None] & k_ok[None, :], other=0.0)
155
- xs = tl.load(xs_ptrs, mask=m_mask[:, None] & s_ok[None, :], other=127)
155
+ xs = tl.load(xs_ptrs, mask=m_mask[:, None] & s_ok[None, :], other=0)
156
w = tl.load(w_ptrs, mask=n_mask[:, None] & k_ok[None, :], other=0.0)
157
- ws = tl.load(ws_ptrs, mask=n_mask[:, None] & s_ok[None, :], other=127)
157
+ ws = tl.load(ws_ptrs, mask=n_mask[:, None] & s_ok[None, :], other=0)
158
acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
159
x_ptrs += BLOCK_K
160
w_ptrs += BLOCK_K
@@ -240,9 +240,9 @@ def _block32_packed_kernel(
240
ws_mask = ws_mask & (wg < K // 32)
241
x = tl.load(x_ptr + rows[:, None] * K + xk, mask=x_mask, other=0.0)
242
w = tl.load(w_ptr + cols[:, None] * K + wk, mask=w_mask, other=0.0)
243
- xs = tl.load(xs_ptr + rows[:, None] * (K // 32) + xg, mask=xs_mask, other=127)
243
+ xs = tl.load(xs_ptr + rows[:, None] * (K // 32) + xg, mask=xs_mask, other=0)
244
ws = tl.load(
245
- ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=127
245
+ ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=0
246
)
247
acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
248
# Keep only products of matching K panels: the block diagonal.