metal : fix idle threads in mul_mv_iq3_xxs for ne00 < 1024

[masterFoad]

Overview

This improves iq3_xxs mul-mv thread utilization when nb32 < 32 and nb32 divides 32.

Each thread currently takes one 32-element chunk based on its thread ID. At ne00 = 512, there are only 16 chunks (nb32 = 16), so threads 16 through 31 have nothing to process. The patch assigns two threads to each chunk, with each thread covering half of the nr0 output rows. When nb32 >= 32, or nb32 does not divide 32, the existing mapping is unchanged.

The PR also changes N_R0_IQ3_XXS from 4 to 8, so each simdgroup handles up to eight rows. This was tuned separately while testing; see the benchmark results below. I can revert it if this feels too specific.

Tested on my 14inch MacBook Pro M5 24 GB , using Tiel-Coder-35B-A3B IQ3_XXS.

On a fixed replay of the same coding-agent workload, the b10488 development tree improved from 24.66 to 22.59 s/rep (−8.4% wall), with decode increasing from 65.6 to 73.9 tok/s and perplexity unchanged at 3.3969.

Additional information

The same (row, chunk) work is still computed once. Threads outside a row's assigned range contribute zero to the final simd_sum, and the full simdgroup reduction produces the intended row result.

test-backend-ops comparisons against the CPU backend pass for MUL_MAT_ID and MUL_MAT, including k = 256, the deepest split, and k = 768, the fallback path.

Performance was measured on upstream/master = daef7b6:

n	master (µs)	row-split only (µs)	PR (split + N_R0 = 8) (µs)	delta vs master
2	143.89	114.25	102.39	−28.8%
4	263.91	218.67	170.11	−35.5%
8	492.38	404.67	311.60	−36.7%
12	710.27	582.84	466.73	−34.3%
16	921.72	738.84	611.68	−33.6%
24	1363.69	1056.90	914.15	−33.0%

The stock test-backend-ops perf keeps the selected expert IDs constant across repetitions, so I used a local benchmark modification that rotates expert IDs between short timed batches to avoid repeatedly benchmarking the same cached weights. I can also submit the rotating-ID benchmark change separately if that would be useful.

The n = 1 case is excluded because it was noisy on my setup.

The measured case is ne00 = 512 (nb32 = 16). Other iq3_xxs matrices with nb32 < 32, where nb32 divides 32, take the same row-split path.

The same ix = tiisg thread assignment appears in six other kernels: iq2_xxs, iq2_xs, iq3_s, iq2_s, iq1_s, and iq1_m. They leave threads idle in the same way when nb32 < 32. I don't have a model that hits those kernels at small ne00, so they're left unchanged here and flagged for follow-up.

Test methodology

• End-to-end benchmark: fixed replay of a multi-turn coding-agent transcript, 13 generations, ~5.6K-token system prompt, warm prefix cache, MTP + ngram speculative decoding.
• Same llama-server configuration for each run; results are the mean of 3 warm repetitions.
• The workload exercises the same speculative verify sizes as the kernel table (n ≈ 2–4 for MTP, up to n ≈ 24 for ngram), with iq3_xxs ffn_down_exps running across 37 blocks on each verify pass.
• The end-to-end numbers above were measured on the b10488 development tree; the kernel table was measured separately against upstream/master = daef7b6.
• Perplexity was checked with llama-perplexity -m Tiel-Coder-35B-A3B-MTP-UD-IQ3_XXS.gguf -f ppl_ref.txt -c 2048 -b 2048 --chunks 4 -ngl 999 -fa on -ctk q8_0 -ctv q8_0; the result remained unchanged at 3.3969.
• Server config:

LLAMA_ARG_TOP_K=1 \
llama-server -m Tiel-Coder-35B-A3B-MTP-UD-IQ3_XXS.gguf \
  --ctx-size 32768 --batch-size 512 --ubatch-size 256 \
  --flash-attn on --cache-type-k q8_0 --cache-type-v q8_0 \
  --parallel 1 \
  --spec-type draft-mtp,ngram-mod \
  --spec-draft-n-max 3 --spec-draft-n-min 1 --spec-draft-p-min 0.8 \
  --spec-ngram-mod-n-match 32 --spec-ngram-mod-n-max 24 --spec-ngram-mod-n-min 8

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, AI was used during the implementation and profiling, guided and reviewed by me.

[forforever73]

For nb32 >= 32, the lane mapping is unchanged, but changing N_R0 from 4 to 8 doubles the rows per simdgroup and regresses performance on my device.

[masterFoad]

I will revert this, and rerun the bench, thanks!
Can you please share test params/results that you got and hardware?

[forforever73]

I ran the following command on M4 Max:

test-backend-ops perf -b MTL0 -o MUL_MAT --test-file tests.txt

The tests.txt contains:

29 0 4096 1 1 1 0 2 18 512 4096 1 1 0 0 0 0 0 512 1 1 1 0 0 0 0 -
29 0 4096 1 1 1 0 2 18 768 4096 1 1 0 0 0 0 0 768 1 1 1 0 0 0 0 -
29 0 4096 1 1 1 0 2 18 1024 4096 1 1 0 0 0 0 0 1024 1 1 1 0 0 0 0 -

and results are:

Parameters	Baseline	N_R0=8	Latest PR
m=4096,n=1,k=512 (nb32=16)	7.49 µs	6.58 µs	6.68 µs
m=4096,n=1,k=768 (nb32=24)	9.09 µs	9.67 µs	9.20 µs
m=4096,n=1,k=1024 (nb32=32)	8.76 µs	9.46 µs	8.68 µs

The performance regression is gone now.

[masterFoad]

Thanks for testing this!
Your M4 Max results match the intended behavior exactly: k=512 uses the split path and improves, while k=768 and k=1024 stay on the original path and remain at baseline.

I also reproduced the regression you found and changed the implementation so the normal kernel stays at N_R0=4; the 8-row version is now only used for the narrow split case.

On my M5, the regression is gone and the ne00=512 case keeps the performance gain. Thanks for catching it.

[ggerganov]

Is this valid only for IQ3? Can we generalize across all types?

[masterFoad]

@ggerganov Thank you for taking a look!

Yes, the idea generalizes:

iq1_s, iq1_m, iq2_xxs, iq2_xs, iq2_s, iq3_xxs, and iq3_s all use the same lane-to-block mapping, so they have the same idle issue when nb32 < 32 (ne00 < 1024), For the narrow shapes where nb32 divides 32, the split approach used here should generalize directly to those kernels.

The K-quants have the same underlying issue, but use different lane mappings, so they would need a separate implementation.

I suggest to extend this PR to the other iq kernels and handle the K-quants in a separate PR. Happy to do the follow up either way, WDYT?

[ggerganov]

Can you try to make the SPLIT a function constant instead of template argument? If the performance is the same, the FC would be preferable.

[masterFoad]

Done, performance is unchanged on my M5.

[ggerganov]

Nice! Looking forward to the generalization of this idea