New tuning for small K gemv

[jagrit06]

Proposed changes

• Add a new tuning for small k gemv
• Add a new tuning for small output, long K gemv

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[jagrit06]

Before:

  B,     M,     N,     K,   dtype,  t,   gpbs_mx, glops_mx
  1,     1,  4096,    64, float16, nt,    25.342,    24.946
  1,     1, 12288,    64, float16, nt,    43.023,    42.358
  1,     1,    64,  4096, float16, nt,    15.712,    15.466
  1,     1,    64, 12288, float16, nt,    23.527,    23.163

After:

  B,     M,     N,     K,   dtype,  t,   gpbs_mx, gflops_mx
  1,     1,  4096,    64, float16, nt,    36.587,    36.016
  1,     1, 12288,    64, float16, nt,    70.071,    68.988
  1,     1,    64,  4096, float16, nt,    29.781,    29.316
  1,     1,    64, 12288, float16, nt,    55.257,    54.403

[awni]

Looks great!

[ivanfioravanti]

A small preview of the effect on mlx-lm 🚀
cat 4k.txt | python -m mlx_lm generate --model mlx-community/Qwen3-30B-A3B-Instruct-2507-4bit -m 200 --temp 0.7 --top-k 20 --top-p 0.8 --prompt -

before
Generation: 200 tokens, 84.179 tokens-per-sec

after
Generation: 200 tokens, 90.274 tokens-per-sec

[awni]

Huh - that's surprising! Are you sure it's from this PR? I don't think it should affect generation for that model unless I am missing something

[ivanfioravanti]

I installed latest from mlx while doing tests and noticed faster performance and I thought was dut to this commit, probably is another change. 🤔

[ivanfioravanti]

You are right @awni #2608 is the game change here!