cuda: bound mmvq rows by the per-slot stride for MUL_MAT_ID - #258
Merged
Merged
Conversation
With ids, stride_col_dst is the token stride, not the row count. When rows_per_cuda_block > 1 (small_k) and the row count is not a multiple of it, the tail block wrote past the end of its expert slot into the next one, racing with that slot's block. test-backend-ops MUL_MAT_ID with ptq1_0, m=70, n=1, k=2048 failed in 8 of 10 runs on an RTX 4090.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Fixes an out-of-bounds write in
mul_mat_vec_qon the MUL_MAT_ID path. The kernel guarded its row writes (and the bias prefetch) withstride_col_dst, which is the row count only for plain MUL_MAT. Withids,stride_col_dstis the token stride (s2), so the guard let rows past the end of an expert slot through.Why
When
small_kputs several rows in one CUDA block and the row count is not a multiple ofrows_per_cuda_block, the last block of each slot wrote its extra rows into the first rows of the next slot, racing with that slot's own block. The result depends on block scheduling, so it shows up as a flaky test rather than a hard failure. On an RTX 4090 atprism0324c66:test-backend-ops -o MUL_MAT_ID -p ptq1_0failed in 8 of 10 runs, always onMUL_MAT_ID(ptq1_0, n_mats=4, n_used=2, m=70, n=1, k=2048), with errors from 0.005 to 0.17.compute-sanitizerinitcheck, racecheck and memcheck are all clean, since the write lands inside the dst buffer.Real models rarely hit it (expert row counts are usually multiples of 64), but it makes the PTQ1_0 MUL_MAT_ID test unreliable.
How
For
ids, the per-slot row stride isstride_channel_dst(s1), which equals the row count for the contiguous dst this path already assumes. The kernel now takesnrows_dst = ids ? stride_channel_dst : stride_col_dstand uses it in both guards. The plain MUL_MAT path is unchanged, andmul_mat_vec_q_moealready usesnrows_x.Upstream
ggml/src/ggml-cuda/mmvq.cuhas the same two guards.Testing
RTX 4090 (sm_89), CUDA 12.8:
test-backend-ops -o MUL_MAT, all types: 1371/1371.test-backend-ops -o MUL_MAT_ID, all types: 1035/1035 in 4 of 4 runs.Not tested: AMD/HIP (same kernel, so same fix applies) and other NVIDIA architectures. I did not re-measure speed after the change; it only swaps the value a guard compares against.
AI usage disclosure: YES. Claude Code was used to find the cause and write the fix. I reviewed the change and take responsibility for it.