Blockscaled GEMV accumulates past K when floor(K/tile) is not a multiple of the stage count: example 91 fails for batch > 1
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
Split out from #3536, which notes that the pipeline can issue loads past the current K
range. That part is not only an out-of-bounds read: it changes the accumulated result,
and it is observable with the example's own verification as soon as batch > 1.
Mechanism
The mainloop consumes a whole buffer per iteration and the FMA sequence inside an
iteration is unconditional:
// include/cutlass/gemm/kernel/gemv_blockscaled.h
int tile_idx = 0;
while (tile_idx < total_tiles) {
int smem_pipe_read_curr = smem_pipe_read;
for_each(make_int_sequence<kStageCount>{}, [&] (auto k_block)
{
...
accum += blockscaled_multiply_add(fragA_reg[frag_idx], fragB_reg[frag_idx],
fragSFA_reg[frag_idx], fragSFB_reg[frag_idx]);
});
tile_idx += kStageCount;
}
Nothing guards tile_idx + k_block against total_tiles, so the loop always performs
ceil(total_tiles / kStageCount) * kStageCount FMAs. Whenever
total_tiles % kStageCount != 0 it accumulates over K beyond gemm_k. The prologue
primes a full buffer with valid_tile = true unconditionally, so the loads backing
those FMAs are issued as well.
For the fp4 configuration in example 91: tileA_k_local = kThreadsPerRow * kElementsPerAccess = 8 * 32 = 256 and kStageCount = 4, so the condition is
floor(K / 256) % 4 != 0.
Reproduction
sm_120a, CUDA 13.0.88, RTX 5070 Ti Laptop, example 91:
./91_fp4_gemv --m=256 --k=1280 --batch=2 --epilogue_st=1.0 --profiling=false
fails verification. The same K at --batch=1 passes.
Sweep, all with --m=256 --epilogue_st=1.0 --profiling=false:
| K | floor(K/256) % 4 | batch=1 | batch=2 | batch=3 |
|---|---|---|---|---|
| 256, 288, 512, 544, 768, 800 | != 0 | pass | fail | fail |
| 1280, 1312, 1536, 1568, 1792 | != 0 | pass | fail | fail |
| 2304 | != 0 | pass | fail | fail |
| 1024, 1056, 1152, 2048, 2176 | == 0 | pass | pass | pass |
At batch = 1 the reads past gemm_k do not change the verified output for the
allocations the example makes; I have not established why, so I would not rely on
that. At batch > 1 they reach the following batch's operands and the result is wrong.
The table above was measured with the K-tail fix from #3536 applied (#3567). The same
batch > 1 failures are present without it, so this is independent of that fix.
Scope
This is distinct from the dropped K tail reported in #3536 and is not addressed by
#3567. Fixing it looks like it needs per-stage load/FMA predication together with a
matching change to how the trailing partial stage is accumulated, rather than a change
to the tail path.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in include/cutlass/gemm/kernel/gemv_blockscaled.h, focusing on the mainloop, prologue, and stage indexing shown in the issue. Reproduce with example 91 using --m=256 --k=1280 --batch=2 --epilogue_st=1.0 --profiling=false, then verify the K sweep passes for batches 1, 2, and 3 without accumulating beyond gemm_k.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- hpc, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100