NVIDIA / NVIDIA/cutlass

Blockscaled GEMV accumulates past K when floor(K/tile) is not a multiple of the stage count: example 91 fails for batch > 1

Open
#3,568 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

CUTLASS C++
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.