Skip to content

[CUDA] FPxINT quantized matmul for Hopper - #3160

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:qmm-sm90
Feb 25, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:qmm-sm90

Conversation

@zcbenz

@zcbenz zcbenz commented Feb 23, 2026 •

Copy link
Copy Markdown
Member

Refs #2536.

This PR implements QMM for Hopper only, using CUTLASS.

Performance of FP16xINT4 with group size 64 on H100:

M N K CUTLASS TFlops CUBLAS TFlops Speedup
1 16384 16384 6.0 2.8 2.16
2 16384 16384 12.0 5.6 2.15
4 16384 16384 24.0 11.4 2.10
8 16384 16384 48.0 21.8 2.20
16 16384 16384 96.1 47.8 2.01
32 16384 16384 100.1 86.9 1.15
64 16384 16384 102.7 178.2 0.58

Independent C++ code for profiling the kernel

For small batch size (<=16) the performance data is similar to (though slower than) Machete so I think it is doing good. When M becomes larger than 32 it can no longer surpass cuBLAS, and it can be solved by increasing tile size, which I will fix in following PRs.

Another problem is this kernel behaves badly for small N/K (0.5x for n4096k4096), I'm not sure if it is something TMA kernels are not good at and we might have to fallback to Ampere optimized kernels.


There are also a few things need further optimizations:

  1. The CUTLASS kernel requires scales/zeros to be N-major but MLX's quant format is K-major so there is a transpose involved for each call. CUDA graph can partly hide this overhead but I haven't profiled how much it hurts performance.
  2. By shuffling the quantized weights we can use wider vectorized reads and faster int-to-float conversions, which would bring 17%~20% improvement. The 55_hopper_int4_bf16_gemm example implements a --shuffle switch that can be used to measure the performance difference.

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me.

The performance seems a bit low ~2.1x vs >3x expected. Do you have a guess as to why?

@zcbenz

zcbenz commented Feb 24, 2026

Copy link
Copy Markdown
Member Author

Skimming the Machete code which achieve 3x for batch size 16/32, I think the performance difference mainly comes from a better designed weights layout. (Machete is modified from the cutlass kernel this PR uses.)

@jagrit06 jagrit06 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should merge this! That said, for the sizes of M we have, we should also be reporting the memory bandwidths since flops isn't the metric we will be maxing out

@zcbenz
zcbenz merged commit 6304c28 into ml-explore:main Feb 25, 2026
16 checks passed
@zcbenz
zcbenz deleted the qmm-sm90 branch February 25, 2026 00:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants