Skip to content

cuda: bound mmvq rows by the per-slot stride for MUL_MAT_ID - #258

Merged
bri-prism merged 1 commit into
prismfrom
fix/mmvq-ids-row-guard
Sep 29, 2026
Merged

bri-prism merged 1 commit into
prismfrom
fix/mmvq-ids-row-guard

Conversation

@bri-prism

Copy link
Copy Markdown
Collaborator

What

Fixes an out-of-bounds write in mul_mat_vec_q on the MUL_MAT_ID path. The kernel guarded its row writes (and the bias prefetch) with stride_col_dst, which is the row count only for plain MUL_MAT. With ids, stride_col_dst is the token stride (s2), so the guard let rows past the end of an expert slot through.

Why

When small_k puts several rows in one CUDA block and the row count is not a multiple of rows_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 at prism 0324c66:

  • test-backend-ops -o MUL_MAT_ID -p ptq1_0 failed in 8 of 10 runs, always on MUL_MAT_ID(ptq1_0, n_mats=4, n_used=2, m=70, n=1, k=2048), with errors from 0.005 to 0.17.
  • A shape probe showed it is not PTQ1_0 specific: PQ2_0 fails the same way at m=70 with k=1024 and 2048. m=64, 72, 128 and 200 pass, as does plain MUL_MAT at m=70.
  • compute-sanitizer initcheck, racecheck and memcheck are all clean, since the write lands inside the dst buffer.
  • It is older than cuda: branch-free PTQ1_0 MMQ tile loader and full Ampere tile table (2x prefill) #214: the same case fails at bdc23b5.

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 is stride_channel_dst (s1), which equals the row count for the contiguous dst this path already assumes. The kernel now takes nrows_dst = ids ? stride_channel_dst : stride_col_dst and uses it in both guards. The plain MUL_MAT path is unchanged, and mul_mat_vec_q_moe already uses nrows_x.

Upstream ggml/src/ggml-cuda/mmvq.cu has the same two guards.

Testing

RTX 4090 (sm_89), CUDA 12.8:

  • The shape probe above (PTQ1_0 and PQ2_0, m in {64, 70, 72, 128, 200}, k in {1024, 2048, 4096}, MUL_MAT and MUL_MAT_ID): 6 of 6 runs clean, previously every m=70 MUL_MAT_ID case failed.
  • 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.

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.
@bri-prism
bri-prism merged commit 88c4bc6 into prism Sep 29, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant