Skip to content

Fix dropped K tail in blockscaled GEMV - #3567

Open
TANGBUDU wants to merge 1 commit into
NVIDIA:mainfrom
TANGBUDU:fix/blockscaled-gemv-k-tail
Open

Fix dropped K tail in blockscaled GEMV#3567
TANGBUDU wants to merge 1 commit into
NVIDIA:mainfrom
TANGBUDU:fix/blockscaled-gemv-k-tail

Conversation

@TANGBUDU

@TANGBUDU TANGBUDU commented Aug 29, 2026

Copy link
Copy Markdown

Addresses #3536.

Problem

gemv_blockscaled.h gates the K tail on unroll_col_k, which is the gmem prefetch
issue cursor rather than the K range that has been accumulated. The prologue primes
one buffer and every mainloop iteration issues one more ahead of consumption, so at
loop exit unroll_col_k leads the accumulated range by kStageCount * tileA_k_local.

When floor(gemm_k / tileA_k_local) is a multiple of kStageCount and gemm_k is
not tile aligned, the guard is false for every thread, process_tail_elements() never
runs, and up to a tile of K is dropped from the dot product.

Fix

tile_idx advances by kStageCount per iteration and the FMA sequence within an
iteration is unconditional, so tile_idx * tileA_k_local at loop exit is the K range
that was accumulated. It is used as the tail origin for both the guard and
process_tail_elements(). ptr_A and ptr_B are not advanced by the mainloop, so
absolute-K addressing inside the tail is unaffected.

Tests

ctest cases registered on example 91, all with
--m=256 --batch=1 --epilogue_st=1.0 --profiling=false:

K floor(K/256) before after
288 1 pass pass
1024 4 pass pass
1056 4 fail pass
1152 4 fail pass
1280 5 pass pass
1312 5 pass pass
2048 8 pass pass
2176 8 fail pass

288, 1280 and 1312 are controls: they would fail if total_tiles * tileA_k_local were
used as the tail origin, which double counts what the mainloop already consumed.

Wider sweep, 21 values of K from 32 to 2304 across batch {1, 2, 3}, same build with
only the kernel header swapped: 29/63 pass before, 38/63 after. The nine that change
are K = 1056, 1152 and 2176 at each batch count. No regressions.

Built and run on sm_120a (RTX 5070 Ti Laptop), CUDA 13.0.88.

Not covered here

#3536 also notes loads issued past the current K range. Because the per-iteration FMA
sequence is unconditional, the mainloop always runs a whole multiple of kStageCount
tiles and keeps accumulating past gemm_k whenever
floor(gemm_k / tileA_k_local) % kStageCount != 0. At batch=1 those configurations
happen to verify; at batch > 1 they do not:

./91_fp4_gemv --m=256 --k=1280 --batch=2 --epilogue_st=1.0 --profiling=false

Filed separately as #3568. Fixing it needs per-stage load/FMA predication plus a
matching change to the partial-stage accumulation, which is a much larger change than
this one, so it is left out here. I can follow up with it if that is useful.

Thanks to @VaggelisGian for the report and root-cause analysis.

The tail guard used unroll_col_k, which is the gmem prefetch issue cursor. The
prologue primes one buffer and each mainloop iteration issues one more ahead of
consumption, so at loop exit unroll_col_k leads the accumulated K range by
kStageCount * tileA_k_local.

When floor(gemm_k / tileA_k_local) is a multiple of kStageCount and gemm_k is
not tile aligned, the guard is false for every thread, process_tail_elements()
is skipped, and up to a tile of K is dropped from the dot product.

tile_idx advances by kStageCount per iteration and the FMA sequence within an
iteration is unconditional, so tile_idx * tileA_k_local at loop exit is the K
range that was accumulated. Use it as the tail origin. ptr_A and ptr_B are not
advanced by the mainloop, so the absolute-K addressing inside
process_tail_elements() is unaffected.

ctest cases on example 91: k=1056, 1152 and 2176 fail before this change and
pass after. k=288, 1280 and 1312 are controls against using
total_tiles * tileA_k_local as the origin, which would double count.

Addresses NVIDIA#3536, reported by @VaggelisGian.

Not covered here: the same issue notes loads issued past the current K range.
Since the FMA sequence is unconditional, the mainloop always runs a whole
multiple of kStageCount tiles and accumulates past gemm_k when
floor(gemm_k / tileA_k_local) % kStageCount != 0. Visible as wrong results for
batch > 1, e.g. --m=256 --k=1280 --batch=2, and needs per-stage load/FMA
predication.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant