vulkan: Coalesce Q4_K/Q5_K scale loads in mul_mm

[TheBlueMatt]

Some SPIR-V compilers (notably mesa) don't handle the current
vulkan Q4_K/Q5_K scale load pattern in mul_mat particularly well.
While reading three `u8`s from the 12-byte scale array should (at
least on some hardware) result in loading the full 12 bytes in a
single LOAD followed by whatever extraction is needed, at least
the ANV Intel driver really can't practically perform this
optimization.

`mesa`'s unsigned upper bound logic doesn't handle tracking bounds
through ternary, resulting in the `(is < 4) ? ... : is - 4` having
an infinite upper bound (as it cannot prove `is - 4` doesn't
underflow). While this could still be rectified if mesa looked at
the array bounds, it currently doesn't and `glslc` currently emits
SPIR-V that doesn't allow for this optimization anyway (though
maybe it will at some point, see
https://github.com/KhronosGroup/glslang/issues/4206).

In mul_mat_vecq we took a different approach to loading the same
fields. We read the first two bytes we needed from `scale` then
took a branch before deciding whether we needed to read a third
byte. In mesa this did, indeed, lead to a top-level branch with
conditional loads. As such these loads ended up not being
coalesced either (at least in the ANV driver) resulting in
additional instructions in our hot loop.

Instead, here, we go ahead and force loading the full 12 bytes and
extract the bits we need from the packed-u32s instead. In mul_mat
there's a few less ternaries and only one extra shift, so even on
drivers that did optimize the previous loads properly the only
material change should be pulling a few extra bytes into registers
(which on most hardware won't cost anything anyway, though
ironically on Intel it theoretically could). In mul_mat_vecq this
requires a bit of extra math and may read bytes from the u32 that
weren't needed, but it seems likely avoiding the branch is a win
on most platforms.

On Intel Xe2/mesa 26.0.4 with the optimizations from
https://gitlab.freedesktop.org/mesa/mesa/-/work_items/15162,

for shader matmul_id_subgroup_q4_k_f32_f16acc_aligned_l:
 * Instruction Count: 2753 -> 2688
 * SEND Count: 269 -> 261
 * Cycle Count: 273976 -> 266138
 * Max live registers: 248 -> 246
 * Non SSA regs after NIR: 381 -> 382

for shader matmul_id_subgroup_q5_k_f32_f16acc_aligned_l:
 * Instruction Count: 2767 -> 2702
 * SEND Count: 271 -> 263
 * Cycle Count: 274140 -> 268144
 * Max live registers: 248 -> 246
 * Non SSA regs after NIR: 381 -> 382

for shader mul_mat_vec_id_q4_k_q8_1_f32:
 * Instruction Count: 1930 -> 1646
 * SEND Count: 116 -> 71
 * Cycle Count: 1348306 -> 843350
 * Max live registers: 78 -> 84
 * Non SSA regs after NIR: 300 -> 135

for shader mul_mat_vec_id_q5_k_q8_1_f32:
 * Instruction Count: 2207 -> 1922
 * SEND Count: 131 -> 86
 * Cycle Count: 1392012 -> 1037836
 * Max live registers: 90 -> 90
 * Non SSA regs after NIR: 300 -> 135

for shader mul_mat_vec_q4_k_q8_1_f32:
 * Instruction Count: 2029 -> 1749
 * SEND Count: 111 -> 66
 * Cycle Count: 1347278 -> 840118
 * Max live registers: 74 -> 80
 * Non SSA regs after NIR: 299 -> 134

for shader mul_mat_vec_q5_k_q8_1_f32:
 * Instruction Count: 2307 -> 2022
 * SEND Count: 126 -> 81
 * Cycle Count: 1379820 -> 954042
 * Max live registers: 86 -> 86
 * Non SSA regs after NIR: 299 -> 134

On one Arc Pro B60, unsloth/Qwen3.5-35B-A3B-GGUF:UD-Q4_K_XL:
 * pp512: 907.34 ± 9.28 -> 941.94 ± 10.53 (+4%)
 * pp2048: 897.95 ± 1.82 -> 931.55 ± 1.79 (+4%)
 * tg128: 49.49 ± 0.02 -> 49.86 ± 0.05 (+ <1%)

On one Arc Pro B60, unsloth/Qwen3.5-27B-GGUF:Q4_K_S:
 * pp512: 324.13 ± 10.52 -> 354.33 ± 6.81 (+9%)
 * pp2048: 329.80 ± 0.25 -> 357.10 ± 0.06 (+8%)
 * tg128: 17.11 ± 0.01 -> 18.11 ± 0.01 (+6%)

On four Arc Pro B60s, unsloth/Qwen3.5-122B-A10B-GGUF:Q5_K_S with
-sm layer (note that -sm tensor improvements will naturally be
less):
 * pp512: 264.55 ± 2.81 -> 280.45 ± 3.94 (+6%)
 * pp2048: 319.32 ± 2.72 -> 335.70 ± 3.48 (+5%)
 * tg128: 26.39 ± 0.01 -> 26.67 ± 0.01 (+1%)

[jeffbolznv]

I didn't double check all the index calculations, but I'm fine with this in theory if it's generally positive across Intel/AMD hardware (this path isn't as important for NV hardware, so no objection if it gets a bit slower).

[TheBlueMatt]

Oops, duh, of course the same load tweak can be applied to mul_mat_vecq to get a speedup on tg. Given that's the hot loop on tg might be worth benchmarking elsewhere. Given the branch was removed there may be a tiny bit more memory traffic but I have to imagine it won't be a loss given the whole cacheline is gonna be read one way or another.

[0cc4m]

Sorry @TheBlueMatt I overlooked this one. I'll test and review it soon.

[0cc4m]

I see good improvements across my hardware:

AMD Radeon RX 9070 XT:

model	size	params	ngl	fa	test	t/s (before)	t/s (after)	diff
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	pp512	3954.47 ± 13.34	4158.75 ± 5.00	+5.2%
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	tg128	115.03 ± 0.06	116.39 ± 0.12	+1.2%
llama 13B Q5_K - Medium	8.60 GiB	13.02 B	99	1	pp512	2073.20 ± 1.53	2116.85 ± 1.18	+2.1%
llama 13B Q5_K - Medium	8.60 GiB	13.02 B	99	1	tg128	60.01 ± 0.01	60.46 ± 0.01	+0.7%

AMD Radeon 8060S:

model	size	params	ngl	fa	test	t/s (before)	t/s (after)	diff
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	pp512	1083.41 ± 12.48	1122.02 ± 16.27	+3.6%
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	tg128	41.57 ± 0.48	43.75 ± 0.39	+5.2%
llama 13B Q5_K - Medium	8.60 GiB	13.02 B	99	1	pp512	499.52 ± 23.80	507.58 ± 26.99	+1.6%
llama 13B Q5_K - Medium	8.60 GiB	13.02 B	99	1	tg128	22.05 ± 0.05	22.48 ± 0.03	+2.0%

AMD Radeon Pro VII:

model	size	params	ngl	fa	test	t/s (before)	t/s (after)	diff
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	pp512	705.43 ± 1.28	704.41 ± 1.07	-0.1%
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	tg128	82.56 ± 0.16	97.04 ± 0.23	+17.5%
llama 13B Q5_K - Small	15.18 GiB	23.57 B	99	1	pp512	187.99 ± 0.52	186.96 ± 0.47	-0.5%
llama 13B Q5_K - Small	15.18 GiB	23.57 B	99	1	tg128	25.41 ± 0.04	30.69 ± 0.12	+20.8%

Intel A770:

model	size	params	ngl	fa	test	t/s (before)	t/s (after)	diff
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	pp512	788.33 ± 1.12	787.77 ± 1.50	-0.1%
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	tg128	47.24 ± 0.04	50.66 ± 0.04	+7.2%
llama 13B Q5_K - Small	15.18 GiB	23.57 B	99	1	pp512	205.79 ± 0.11	205.81 ± 0.11	+0.0%
llama 13B Q5_K - Small	15.18 GiB	23.57 B	99	1	tg128	16.32 ± 0.00	18.10 ± 0.01	+10.9%

Nvidia RTX 3090 (coopmat1):

model	size	params	ngl	fa	test	t/s (before)	t/s (after)	diff
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	pp512	3526.03 ± 14.52	3571.15 ± 40.10	+1.3%
llama 8B Q4_K - Small	4.36 GiB	8.03 B	99	1	tg128	127.56 ± 0.29	127.85 ± 1.10	+0.2%
llama 13B Q5_K - Small	15.18 GiB	23.57 B	99	1	pp512	889.22 ± 4.47	906.52 ± 6.42	+1.9%
llama 13B Q5_K - Small	15.18 GiB	23.57 B	99	1	tg128	42.71 ± 0.12	47.62 ± 0.12	+11.5%

[TheBlueMatt]

@jeffbolznv any chance you mind taking another look here? @0cc4m indicated this is a nontrivial win in some cases for NV now.

[0cc4m]

Ah sorry, this isn't waiting for Jeff necessarily, just for another approval. @ggml-org/maintainers