Skip to content

[Metal] Tile fp quantized_matmul for 6-16 rows - #4712

Open
dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:qmm-small-rows
Open

dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:qmm-small-rows

Conversation

@dhiltgen

Copy link
Copy Markdown
Contributor

This change speeds up speculative and batched decoding fp quant models when generating 6-16 sequences. Main today either reads all weights once per row, or Split-K qmm pads to 32-row tiles and wastes most of its work. This optimization adds a qmv_mma (M1-M4) using 8x8 simdgroup matrices and qmv_nax, which read each block of weights once, unpack it and multiply the rows on the matrix units to reduce memory traffic. This achieves ~3x speedup on M1/M2 generation, and ~2x on newer systems.

mlx_lm.benchmark -p 512 -g 128 -b <batch> performance examples:

M5 Max

Model Batch Generation tok/s main this branch Generation speedup
Qwen3.8-27B-nvfp4 1 33.6 33.6 1.00x
Qwen3.8-27B-nvfp4 8 113.1 187.9 1.66x
Qwen3.8-27B-nvfp4 16 149.5 298.4 2.00x
Qwen3.8-27B-mxfp4 16 149.7 313.6 2.10x
Qwen3.8-27B-mxfp8 16 128.0 219.8 1.72x

(No observed prompt or batch 1 change)

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with coding agent

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Oct 19, 2026
dhiltgen added a commit to dhiltgen/ollama that referenced this pull request Oct 20, 2026
Quantized matmul tiled for 6-16 rows, the shape of speculative verify,
rebased onto ml-explore/mlx#4709. Applied to the pinned MLX until the
pin includes it; drop this commit then.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants