Skip to content

[Fix] conv3d fp8: guard padded M rows in the epilogue store #969

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
coderfeli merged 2 commits into from
Aug 6, 2026
Merged

[Fix] conv3d fp8: guard padded M rows in the epilogue store #969

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
891bb75
[Fix] conv3d fp8: guard padded M rows in the epilogue store
Aug 5, 2026
f86a007
conv3d fp8: make the padded-row guard compile-time conditional
Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion kernels/conv/conv3d_implicit_fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,9 @@ def compile_conv3d_implicit_fp8(n, c, d, h, width, k, kt, kh, kw, st, sh, sw, pt

grid_m = (npq + TILE_M - 1) // TILE_M
grid_n = (k + TILE_N - 1) // TILE_N
# Padded M rows only exist when npq is not tile-aligned; mirrors _row_chk in
# the BF16 kernel so aligned shapes emit no extra compare.
row_chk = npq % TILE_M != 0
elem_ty = fx.Float8E4M3FN

@flyc.kernel(known_block_size=[BLOCK_THREADS, 1, 1])
Expand Down Expand Up @@ -456,7 +459,13 @@ def store_half_pair(acc0, acc1, m_half):
n_idx = row // dhw
sp = row % dhw
off_ncdhw = n_idx * (k * dhw) + col * dhw + sp
buffer_ops.buffer_store(out.to(fx.BFloat16), y_rsrc, off_ncdhw, mask=col_valid)
# Padded M rows (npq..grid_m*TILE_M) must not store: with the
# n==1 layout col*dhw+row they alias real elements of a later
# column and race legitimate stores from another block.
store_mask = col_valid
if const_expr(row_chk):
store_mask = col_valid & (row < fx.Index(npq))
buffer_ops.buffer_store(out.to(fx.BFloat16), y_rsrc, off_ncdhw, mask=store_mask)

store_half_pair(acc00, acc01, 0)
store_half_pair(acc10, acc11, 1)
Expand Down
Loading