ggml-cuda: better vram to lds loading pipeline in load_tiles_q8_0

[iacopPBK]

Overview

This pr rewrites some portions of the load_tile_q8_0 function, to get a good performance improvement with the MI50.
It splits the loading tiles loop in two phases, from vram to registers, then from registers to LDS, instead of doing both operations in a single loop. The scale loading loop is like upstream. Another thing is that when need_check case happens now it masks to zero the oob data writes, not a big deal but maybe cleaner.

Quick bench results:   

GPU	Model	Test	Before (t/s)	After (t/s)	Δ (t/s)	% Δ
MI50	qwen3.5 4B Q8_0	pp512	991.15	1275.10	+283.95	+28.6%
		pp2048	1026.70	1326.90	+300.20	+29.2%
		pp8192	902.59	1173.84	+271.25	+30.0%
		tg128	78.41	79.33	+0.92	+1.2%
		pp512 @ d8192	802.33	1017.33	+215.00	+26.8%
		pp2048 @ d8192	809.52	1018.06	+208.54	+25.8%
		pp8192 @ d8192	752.68	939.88	+187.20	+24.9%
		tg128 @ d8192	71.00	74.14	+3.14	+4.4%
	qwen3moe 30B.A3B Q8_0	pp512	823.79	1122.05	+298.26	+36.2%
		pp2048	1139.30	1474.63	+335.33	+29.4%
		pp8192	932.11	1148.72	+216.61	+23.2%
		tg128	87.12	87.01	−0.11	−0.1%
		pp512 @ d8192	584.91	715.96	+131.05	+22.4%
		pp2048 @ d8192	709.88	823.69	+113.81	+16.0%
		pp8192 @ d8192	592.83	676.08	+83.25	+14.0%
		tg128 @ d8192	76.71	77.28	+0.57	+0.7%
RX 6800 XT	qwen3.5 4B Q8_0	pp512	2857.69	2860.56	+2.87	+0.1%
		pp2048	2919.32	2922.99	+3.67	+0.1%
		pp8192	2708.28	2713.23	+4.95	+0.2%
		tg128	68.43	69.17	+0.74	+1.1%
		pp512 @ d8192	2431.01	2439.08	+8.07	+0.3%
		pp2048 @ d8192	2386.97	2392.60	+5.63	+0.2%
		pp8192 @ d8192	2252.48	2258.32	+5.84	+0.3%
		tg128 @ d8192	65.77	66.03	+0.26	+0.4%

Mi50 is much faster with this code, while rx6800xt shows no regression.
Tried different models (qwen3/3.5 4b, 30b, mistral) and it works on both gpus, answers are correct at all ctx lenghts.

For old gpus the only concern that comes in my mind is that maybe now i'm using more VGPRS. I looked at the compilation statistics and there is not any net increase in vgprs...i didnt lose occupancy on any kernel variant.

I didn't test the tensor path, no idea if it works, but i think it should behave like upstream (there are no changes in the splitted part, just a variable rename). All the changes are commented in code.

curious to see some benchmarks of both older and newer gpus!

@JohannesGaessler what do you think ?

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES , gfx906 fork was originally wrote with llms, i rewrote this code myself and adjusted for the pr.

[iacopPBK]

Small batches almost the same, huge boost >32

GPU	Model	Test	n_ubatch	Before (t/s)	After (t/s)	Δ (t/s)	% Δ
MI50	qwen3.5 4B Q8_0	pp512	16	372.19	368.14	−4.05	−1.09%
			32	548.11	543.73	−4.38	−0.80%
			64	495.72	704.78	+209.06	+42.17%
			128	736.70	983.22	+246.52	+33.46%
			256	872.75	1138.51	+265.76	+30.46%
			512	988.95	1268.35	+279.40	+28.25%

GPU	Model	Test	n_ubatch	Before (t/s)	After (t/s)	Δ (t/s)	% Δ
RX 6800 XT	qwen3.5 4B Q8_0	pp512	16	670.06	676.87	+6.81	+1.02%
			32	1164.90	1166.36	+1.46	+0.13%
			64	1728.27	1744.54	+16.27	+0.94%
			128	2188.35	2199.30	+10.95	+0.50%
			256	2545.69	2562.66	+16.97	+0.67%
			512	2841.33	2854.67	+13.34	+0.47%

[IMbackK]

Hmm this is wierd. really the compiler should be optimizing this to MUBUF instructions to load directly from global memory into lds, and not use any registers for this at all. Intuatively i would say that your change makes it mutch harder for the compiler to optimize this in. Have you looked at the output instructions, i suspect the compiler is not managing in either version for some reason.

[iacopPBK]

yeah i agree, its like mi50 cant optimize that. I will check before / after instructions for vec dot q8 kernel. dont know if it is the compiler or the gpu itself 😆

[IMbackK]

yeah so i would expect it to use buffer_load_dword with the lds destination parameter, its not doing it in either version

mmq-instance-q8_0-hip-amdgcn-amd-amdhsa-gfx908_pr.txt
mmq-instance-q8_0-hip-amdgcn-amd-amdhsa-gfx908_master.txt

[iacopPBK]

Cool! It must be found a way to hint those direct loads then, and it will be better.

Regarding the oom data masking to 0, do you think is useful or just redundant?

[ORippler]

Do you have some insights how much of the perf gains are due to loop optimization (hoisting invariants outside the loop + replacing if's with ternaries) vs. splitting the loop? Register pressure is quite high, though typically not at loading the tiles into SMEM in MMQ (at least on CUDA).

For old gpus the only concern that comes in my mind is that maybe now i'm using more VGPRS. I looked at the compilation statistics and there is not any net increase in vgprs...i didnt lose occupancy on any kernel variant.

In CUDA, this would indicate that the compiler does not materialize the declared register arrays, but instead merges the loops again, hence my question. Note I have not taken a look at PTX/SASS for the proposed changes.

[iacopPBK]

I'm reading the vega 7nm manual and i'm concerned about the direct load to lsd feasibility:

This part is saying that it truncates the last two bit of the address, so it always rounds to 4byte memory alignment. Q8_0 block is 34bytes so basically never aligned. Maybe this is one of the reasons compiler dont emits buffer dwords to lds.

Another thing that probably the manual is telling me is:

That i can't do the bitwise OR and shitfs implemented in the get_int_b2 because it needs ALU (that require vgprs).

So at this point i'm not sure we are doing any direct load to LDS at all in every GPU with this function, even though i don't know how the cp.async would work here and if it supports alu ops on LDS.

Maybe its just that the vega scheduler can't realize that the instructions can be compacted, while the 6800xt one can.

[iacopPBK]

Do you have some insights how much of the perf gains are due to loop optimization (hoisting invariants outside the loop + replacing if's with ternaries) vs. splitting the loop? Register pressure is quite high, though typically not at loading the tiles into SMEM in MMQ (at least on CUDA).

For old gpus the only concern that comes in my mind is that maybe now i'm using more VGPRS. I looked at the compilation statistics and there is not any net increase in vgprs...i didnt lose occupancy on any kernel variant.

In CUDA, this would indicate that the compiler does not materialize the declared register arrays, but instead merges the loops again, hence my question. Note I have not taken a look at PTX/SASS for the proposed changes.

The speedup is purely related to the split loop. Keeping that stuff outside the loop didn't change anything, it only looked "cleaner" to me. I'm not sure about the ternary vs ifs cause i replaced them directly.

[JohannesGaessler]

As I said before: please take a look at the code I implemented in the tile and vector FA kernels. Many of the same optimizations I implemented there will also apply to MMQ. Even a simple patch like JohannesGaessler@2bb9ae0 outperforms this PR in my testing:

model	n_ubatch	test	t/s master	t/s PR	t/s 16 threads / row
llama 8B Q8_0	16	pp512	116.78	179.07	238.53
llama 8B Q8_0	32	pp512	398.32	394.51	399.41
llama 8B Q8_0	64	pp512	342.03	469.89	519.34
llama 8B Q8_0	128	pp512	372.71	515.81	573.46
llama 8B Q8_0	256	pp512	455.34	612.88	684.63
llama 8B Q8_0	512	pp512	513.45	689.47	770.32

Also the code you've submitted is just too convoluted and complex. I would really not want to merge and maintain it if there is an alternative.

[iacopPBK]

I made a bench swiping the threads_per_row in the new commit:

Its a matter of testing and tuning the best number of threads.
Im gonna read the fattn files and i will update the code with the thread selection. Maybe with a runtime check the best value can be always kept, otherwise just the best overall.

MI50

Test	Threads	16	32	64	128	256	512	1024	2048
pp512	8	399.27	409.26	772.42	1086.54	1251.35	1322.74	1292.65	1267.86
pp512	16	425.70	548.56	766.34	1076.06	1238.42	1314.55	1284.81	1250.65

RX 6800 XT

Test	Threads	16	32	64	128	256	512	1024	2048
pp512	1	660.48	1134.09	1737.10	2264.06	2704.69	3007.41	3010.85	3001.95
pp512	2	709.97	1226.39	1836.33	2273.15	2659.37	2943.49	2876.43	2941.17

[iacopPBK]

I added a threads_per_row selector.

With these threads_per_row i tried to split the loops again as the first commit and there was no speedup. No way to use memcpy here because the loads are 2byte aligned so no vectorization possible.

@JohannesGaessler do you think there is a better way to assign the threads instead of this empirical approach?

If this kind of selection is ok we can expand it for other archs with some tests.

16 is the best number for vega because it keeps the small batches fast and at long ctx (tested pp32184 -d0,32184), 8 and 16 are almost identical, so runtime check can't help.
Also don't know if mma is already optimal (pretty sure it is for sone gpus at least)

[JohannesGaessler]

I don't think there is a better way to do it than to just benchmark different values and pick the best one. Please be aware that there are concurrent efforts #21849 to change how kernel orchestration is done, any changes to the number of threads used for loading should be made to fit that pattern.