FP_QMM proposal

[Blaizzy]

After benchmarking the current activation quantization approach suggested by @awni in #3122, I spent last couple days exploring quantized matrix multiplication (QMM). I started with a cuBLAS-based implementation as a baseline, then built a custom Tensor Core kernel that achieves significantly better results.

Key findings

• AQ(Actvation Quant) + Tensor Core QMM achieves within 6% of AQ-only (main branch) prompt throughput.
• The custom Tensor Core kernel reaches ~50 TFLOPS for quantized matmul, a 2.5x improvement over the ~20 TFLOPS naive approach.
• The cuBLAS QMM was notably slower, halving prompt throughput and adding ~2 GB peak memory.
• mxfp4 now works and cuts memory nearly in half (9.7 GB) and delivers 21% faster generation than mxfp8.

How it works

Instead of dequantizing weights to full precision then doing matmul (two passes), the kernel dequantizes and computes in a single fused pass using Tensor Cores.

Benchmarks

GPU: RTX 6000 Pro Workstation (96GB VRAM)
Model: Ministral-3-14B

mxfp8

Approach	Prompt (tok/s)	Generation (tok/s)	Peak Memory (GB)
AQ only (main branch)	2,458	101	16.3
QMM (cuBLAS)	928	87	18.3
AQ + QMM (cuBLAS)	1,002	91	16.3
QMM (Tensor Core)	1,518	94	16.6
AQ + Tensor Core QMM	2,320	101	16.3

mxfp4

Approach	Prompt (tok/s)	Generation (tok/s)	Peak Memory (GB)
QMM (Tensor Core)	1,703	123	9.7

If these results look promising, happy to clean this up and send a PR.

I would love to have this be my first PR to mlx-core.

cc: @awni @zcbenz

[zcbenz]

Can you run some benchmarks on the qmm op itself?

Something like https://github.com/IST-DASLab/marlin/blob/master/bench.py that:

• checks the speedup against the dense matmul
• checks common problem sizes and batch sizes

[Blaizzy]

Hey Chen

Thanks for the speedy response!

I ran the benchmark, here is what I got:

  
  Quantization: 4-bit, group_size=32

  Model: ideal
  ----------------------------------------------------------------------
  batch    time(ms)     TFLOP/s      GB/s         speedup
  ----------------------------------------------------------------------
  1        0.027        1.24         331.3        0.94      x
  2        0.042        1.59         212.5        0.78      x
  4        0.045        3.00         200.9        0.77      x
  8        0.056        4.79         161.4        0.59      x
  16       0.071        7.58         129.6        0.59      x
  32       0.113        9.52         83.7         0.36      x
  64       0.279        7.70         35.7         0.16      x
  128      0.301        14.26        36.6         0.14      x

  Model: Llama7B
  ----------------------------------------------------------------------
  batch    time(ms)     TFLOP/s      GB/s         speedup
  ----------------------------------------------------------------------
  1        0.166        2.45         650.3        1.44      x
  2        0.304        2.66         354.4        0.94      x
  4        0.341        4.75         316.8        0.82      x
  8        0.377        8.58         287.6        0.72      x
  16       0.562        11.53        195.2        0.48      x
  32       0.885        14.64        126.3        0.30      x
  64       1.525        16.99        76.0         0.22      x
  128      1.679        30.85        74.1         0.20      x

  Model: Llama13B
  ----------------------------------------------------------------------
  batch    time(ms)     TFLOP/s      GB/s         speedup
  ----------------------------------------------------------------------
  1        0.200        3.17         841.9        2.07      x
  2        0.428        2.97         394.6        1.09      x
  4        0.496        5.12         341.3        1.00      x
  8        0.615        8.26         276.2        0.80      x
  16       0.848        11.97        201.8        0.56      x
  32       1.216        16.70        143.0        0.42      x
  64       1.965        20.66        91.1         0.26      x
  128      2.179        37.27        87.0         0.25      x

  Model: Llama70B
  ----------------------------------------------------------------------
  batch    time(ms)     TFLOP/s      GB/s         speedup
  ----------------------------------------------------------------------
  1        0.459        4.24         1126.2       2.68      x
  2        1.061        3.67         487.9        1.21      x
  4        1.181        6.59         438.6        1.08      x
  8        1.452        10.72        357.7        0.88      x
  16       2.297        13.56        227.2        0.56      x
  32       3.341        18.64        157.6        0.40      x
  64       4.471        27.86        120.0        0.29      x
  128      5.244        47.50        106.0        0.26      x

  Model: Ministral14B
  ----------------------------------------------------------------------
  batch    time(ms)     TFLOP/s      GB/s         speedup
  ----------------------------------------------------------------------
  1        0.195        3.34         886.9        2.35      x
  2        0.448        2.90         385.9        1.15      x
  4        0.502        5.18         345.4        1.04      x
  8        0.617        8.43         282.0        0.85      x
  16       0.864        12.05        203.1        0.59      x
  32       1.238        16.80        143.8        0.44      x
  64       1.984        20.98        92.5         0.27      x
  128      2.208        37.69        88.0         0.26      x

  =========================================================================
  =======
  Detailed Single-Layer Benchmarks (4096 x 4096)
  =========================================================================
  =======


  4-bit quantization (mxfp4):
  ------------------------------------------------------------
  M        Quant(µs)    Dense(µs)    Speedup    BW(GB/s)
  ------------------------------------------------------------
  1        30.6         29.6         0.97      x 291.6
  2        40.3         35.7         0.88      x 222.0
  4        50.0         44.2         0.88      x 179.5
  8        55.8         44.5         0.80      x 162.2
  16       62.9         37.6         0.60      x 145.8
  32       107.0        36.7         0.34      x 88.2
  64       269.9        36.7         0.14      x 36.9
  128      308.6        41.9         0.14      x 35.7
  256      301.1        55.6         0.18      x 43.5
  512      303.2        82.0         0.27      x 57.1

  8-bit quantization (mxfp8):
  ------------------------------------------------------------
  M        Quant(µs)    Dense(µs)    Speedup    BW(GB/s)
  ------------------------------------------------------------
  1        22.8         27.4         1.20      x 758.2
  2        39.1         37.8         0.97      x 442.9
  4        46.9         45.1         0.96      x 370.6
  8        53.6         42.3         0.79      x 325.1
  16       65.5         37.0         0.56      x 268.0
  32       102.6        43.7         0.43      x 173.8
  64       299.0        43.4         0.15      x 61.4
  128      330.4        49.2         0.15      x 58.7
  256      329.1        62.4         0.19      x 65.3
  512      323.7        80.6         0.25      x 79.4

[zcbenz]

Thanks for the running the benchmark, for quantized matmul we expect it to run faster than dense matmul with same shape, otherwise running a quantized model would be slower than running the original weights, especially for the small batch size which is bound by memory bandwidth. The ideal speedup would depend on how much memory bandwidth saved by quantization, for example a FP4xFP16 qmm should be ~4x faster than FP16xFP16 dense matmul.

[Blaizzy]

Yes, agreed!

This morning, I got it from 2x to around 3x. Previously, it was defaulting to the fq_qmv for vectors and not using the new fq_qmm path. Plus, a couple improvements.

MXFP4

Prompt(900 toks): From 1600 tok/s to 1800 tok/s
Generation (100 toks): From 122 tok/s to 145 tok/s 

Now the performance is around 60% the max bandwidth of my card (1.8TB/s) and close to m3 ultra.

My goal is to get it to around 70%-80% which would be 4x :)

I will send you updated results when I get home

[Blaizzy]

Ohh, I see you have a PR already 😁

Thanks a lot @zcbenz, this was fun!

Let me know if there anything else I can help with or test on my card.