Quantized matmul: opt-in mode where a batched row matches M=1

Open
  • Is something wrong?
#4,819 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
35/100
Issue type
Feature
Clarity
Clearly specified
Activity status
Active
Tech stack
cpp, machine-learning
Domain
machine-learning

Research direction

Start in mlx/backend/metal/quantized.cpp at QuantizedMatmul::eval_gpu, use_qmv_wide, and qmm_nax_k_partitions, since they control M-based dispatch and split-K selection. Trace the M=1 and batched quantized-matmul paths, then reproduce the issue by comparing each unequal batched row with the same row run alone. Done means an opt-in flag can force the row-exact route while leaving the default fast path unchanged.

Written by the indexing model from the issue text.

Description

  • I understand it is strictly prohibited to use AI to write issues.

AI usage disclosure: drafted with AI assistance from measurements we already had. The line numbers below are from current main.

Could quantized matmul grow an opt-in mode where a row inside a batch matches that same row run at M=1?

The Metal dispatch picks the kernel from M.

  • use_qmv_wide (mlx/backend/metal/quantized.cpp, around line 543) runs when M is at least 2. M=1 stays on qmv.
  • QuantizedMatmul::eval_gpu compares M with get_qmv_batch_limit (around line 2183) and sends larger M through qmv_mma, qmm_splitk_nax, or qmm_splitk.
  • qmm_nax_k_partitions (around line 1290) sets the split count from M, because the tile count is ceildiv(M, 64).

A different kernel is a different reduction order, so the result changes with the batch.

What we measured: 16 unequal rows in one forward, against the same rows run alone, flipped 30 to 38 of 256 greedy tokens. That was Qwen3-30B-A3B and a 27B. Grouped short prefill had 12 of 126 rows differ.

oMLX carries a hand patch for the row-exact route: omlx/patches/row_exact_qmv.py. vLLM documents the same property as a feature, docs/features/batch_invariance.md, behind VLLM_BATCH_INVARIANT.

The change we run locally turns qmv_wide and split-K off for every call. That costs throughput, so it has to be a flag. Making it the default would slow the fast path for everyone.

Augustas from Malibu

Dominant language
C++
Stars
29.1k
Forks
2.4k
Avg merge
3d 13h
Merged PRs (30d)
108

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from ml-explore/mlx

All issues in ml-explore/mlx

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.