CUDA: fix MMQ out-of-bounds reads

[JohannesGaessler]

Fixup to #29941 . See comments starting at #27044 (comment) .

I was unable to reproduce further issues locally but looking at the way I implemented the utility function ggml_cuda_mmq_get_J_max it is definitely wrong on master in a way that is consistent with the reported issues. On master the number of columns that are being checked for viable configs is being rounded down for the allocation but rounded up when selecting the actual config to run.

The lowest-effort fix would be to modify ggml_cuda_mmq_get_J_max so that the rounding is consistent. However, the fact that the rounding can be inconsistent in the first place is a design flaw. So this PR instead moves the code for determining the optimal tile size forward and directly uses that to calculate how much padding is needed. This also plays nicer with recent efforts to modify the logic for selecting the optimal tile size for MoE models.

The way fallback is being calculated on master can in principle also be inconsistent so I deduplicated it as well.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: No

[hclsys]

Checked on a GB10 (sm_121a, CUDA 13.0). src1_q8_1 in its own exact-size cudaMalloc (the pool alloc swapped for a cudaMalloc of the same size), compute-sanitizer memcheck, cases from the #27044 thread.

Plain memcheck, q4_0, 512 experts, 10 used, m=640, n=100: master 1194 errors (b=0) / 274 (b=1), this PR 0 / 0. n=508 and the q4_K shapes are 0 on both. MUL_MAT + MUL_MAT_ID on CUDA0: 2243/2243.

With --padding 65536 it's still not clean on some routings: reads up to 112 B past the end of src1_q8_1 (size 2896128, your formula), in mul_mat_q<q4_0,112>. The tile_y load loops in mul_mat_q_process_tile run to the next multiple of nthreads (J=112: 4032 ints -> 4096 with 256 threads) and mmq_get_nbytes_shared pads nbs_y for that, but the global padding is only J_bestsizeof(block_q8_1_mmq). With a fixed seed for the ids shuffle in test-backend-ops, 10 of 16 seeds hit it at n=100, and 0 of 16 with GGML_PAD(J_bestsizeof(block_q8_1_mmq), config.nthreads*sizeof(int)) as the padding.

ids_dst gets read up to J_best-1 ints past ne_get_rows too with exact-size allocs; padding it by J_best clears that.

[JohannesGaessler]

Thank you for checking. I agree with those fixes and pushed a corresponding patch, please re-test.

[hclsys]

Re-tested 20c3140 on the GB10, same harness (src1_q8_1, ids_src1, ids_dst at your ne_get_rows + J_best - 1, and expert_bounds each in an exact-size cudaMalloc, memcheck with --padding 65536). The 8 cases from before are 0 errors on 3 seeds, and q4_0 n=100 is 0 on 16/16 fixed ids-shuffle seeds (5bd8b00 failed 10/16).

Also swept q4_0/q4_K/q6_K/q8_0 x n in {9, 33, 65, 129, 257, 400} x m in {512, 576} (128 experts, 8 used), 2 seeds: 96/96 clean. MUL_MAT + MUL_MAT_ID on CUDA0: 2291/2291.

[glennneuber]

@JohannesGaessler , @hclsys thanks for your work. Much appreciated.

The over night agent came to the same conclusion as your corresponding patch.

Checked #29953 with your corresponding patch, all passing:

Hardware	Result
RTX PRO 6000 (sm_120a)	190/190 pass
RTX 2080 Ti (sm_75)	18/18 pass
NVFP4, sm_120a	all pass
Radeon 8060S (RDNA 3.5, gfx1151)	72/72 + 36/36 pass
CPU sweep, 10 architectures	0 shapes short

All GPU runs used exact-size or guard-page allocations

Since it is an MMQ padding issue, also tested on AMD ROCm platform (ggfx1151) and nvfp4 (sm_120->sma_120a).

A potential test is proposed in mmq-rules-check.cu:

How it works
The CPU check replays llama.cpp's tile selection against its real config tables, for every architecture, quant type, batch size and expert count, about 9.6 million shapes with no GPU. For each shape it tests whether the padding covers what the kernel actually reads, which is how it showed the published first #29953 short in 3.5 million shapes and the head short in none.

What it covers
It also checks every architecture, including the ones nobody here has hardware for. That prevents intermittent out-of-bounds reads that crash on users' machines and pass every normal GPU test.

Limitations
It checks the padding rules as modelled, not the compiled code. The guard mmq-tests.patch would cover the code on the respective platforms.
Still untested: MUSA, and AMD CDNA3/RDNA4/GCN hardware, which we only modelled.

Links:

Test cases: https://github.com/MaxusAI/ollama/blob/5273e495fef54972f5d265ee49699e3e39ff1e0d/docs/maxusai/tasks/mmq-tests.patch
Debug allocator: https://github.com/MaxusAI/ollama/blob/5273e495fef54972f5d265ee49699e3e39ff1e0d/docs/maxusai/tasks/mmq-debug-alloc.cuh
CPU padding sweep: https://github.com/MaxusAI/ollama/blob/5273e495fef54972f5d265ee49699e3e39ff1e0d/docs/maxusai/tasks/mmq-rules-check.cu
Write-up: https://github.com/MaxusAI/ollama/blob/5273e495fef54972f5d265ee49699e3e39ff1e0d/docs/maxusai/upstream-mmq-29953-material.md

Disclosure
All the linked code (patches, sweep, allocator, scripts) is AI-generated.

Hope this helps.

[manunicholasjacob]

I re-tested 3070d927f on an RTX 3050 Laptop (sm_86, CUDA 13.2) using an exact-size allocation pool and device-side bounds counters, with the sizes taken directly from the PR's allocations; compute-sanitizer isn't usable under WSL2 here. 3070d927f shows 0 B past both src1_q8_1 and ids_dst across all tested shapes, including the J=32 suite case where 5bd8b0013 over-read in 17/20 runs, and all 406 MMQ ids_dst launches where the earlier build over-read in 381/406. MUL_MAT_ID passes 938/938 and MUL_MAT passes 1304/1304, including with the diagnostics disabled. One additional observation is that the same loads still go past the logical end when measured against the old unpadded allocation sizes, but on 3070d927f they now land in the padding; I did not check whether those values are actually used.

RTX 3050 Laptop (sm_86), CUDA 13.2, exact-size pool + device-side bounds counters, max bytes past src1_q8_1 / ids_dst
| case                                               | master dd266785c      | #29953 5bd8b0013      | #29953 3070d927f |
| q4_0 n=100, b=0 and b=1 (J=112)                    | 2,268 / 436           | 0 / 432               | 0 / 0 |
| q4_0 n=508, b=0 and b=1 (J=128)                    | 0 / 480               | 0 / 476               | 0 / 0 |
| q4_K n=65, fallback (J=128)                        | 8,924 / 500           | 0 / 492               | 0 / 0 |
| q4_K n=100 (J=112)                                 | 2,124 / 432           | 0 / 428               | 0 / 0 |
| q4_K n=2040 (J=128)                                | 0 / 260               | 0 / 148               | 0 / 0 |
| q4_0 + mxfp4, 32 exp, 2 used, n=32 (J=32), 20 runs | 15/20 over, 364 / 120 | 17/20 over, 364 / 120 | 0/20 |
| MUL_MAT_ID suite, MMQ ids launches                 | 9,068 / 504 (381/406) | 364 / 500 (381/406)   | 0 / 0 (0/406) |
MUL_MAT_ID 938/938, MUL_MAT 1304/1304 on all three (also with diagnostics off on 3070d927f)

[JohannesGaessler]

@ynankani as a fixup to #24364 , can you please take care of this?

[ynankani]

Sure, I can take care of this.