support W4A8 Marlin kernel

[HandH1998]

Summary

We inroduce a mixed precision GEMM kernel for INT4-Weight and INT8-Activation. We implemented the W4A8 GEMM based on Marlin GEMM. The kernel is designed to support our W4A8 quantization method QQQ. For more details on the kernel implementation, you can refer to our paper. The kernel demonstrates excellent performance and has been merged into the official vLLM project (see vllm-project/vllm#5218).

We hope the w4a8 GEMM can also provide a practical speedup for other W4A8 quantization methods in the community.
Additionally, since torchao is widely used in frameworks like SGLang, we can extend support for W4A8 once the kernel is integrated into torchao.

Performance

Here is the speedup over PyTorch FP16 GEMM (Calling CUTLASS) of all GEMMs under different numbers of input tokens. The weight matrix size is (N=8192, K=21760). You can reproduce the benchmark results using bench_w4a8.py in my repo.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1113

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• GLIBC not found in Nova workflows

✅ No Failures

As of commit 2690ff4 with merge base 39f16f4 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[drisspg]

can we do some comparisons between this and #880?

[jerryzh168]

thanks, looks pretty good, can you add the benchmark code in https://github.com/pytorch/ao/tree/main/benchmarks as well

and what GPU are the kernels benchmarked on?

[HandH1998]

@drisspg @jerryzh168 I am working on the benckmark and I will give the comparisons with #880. The kernel is benchmarked on A100-80G GPU and it can work on SM > 8.0.

[HandH1998]

@jerryzh168 @drisspg @msaroufim
I have made the following modifications (code modifications refer to #621 and #880):

1. Added benchmark code for marlin_qqq_w4a8 GEMM inbenchmarks/benchmark_marlin_qqq.py
2. Summarized the main differences between marlin_qqq_w4a8 GEMM and marlin_w4a16 GEMM intorchao/quantization/marlin_qqq/README.md
3. Supported marlin_qqq in torchao/quantization/quant_api.py
4. Added some unit tests in test/quantization/test_marlin_qqq.py

w4a8-cutlass is a great work. In comparison, we believe marlin_qqq_w4a8 can support weight per-group quantization in addition to weight per-channel quantization. However, marlin_qqq_w4a8 does have some limitations: it only supports symmetric quantization and the output dtype can only be torch.float16.

In addition, we also provide the performance of torchao/_models/llama/generate.py here. -g128 means weight per-group quantization and the group size is 128.

-q parameter	Precison	Average tokens/sec	Average Bandwidth in GB/s	Peak Memory Usage in GB	Model Size in GB
--compile	fp16	112.45	1486.00	13.93	13.21
-q marlin_qqq --compile	w4a8	197.45	653.50	4.79	3.31
-q marlin_qqq --compile	w4a8-g128	187.62	640.32	4.82	3.41

[HandH1998]

@jerryzh168 @msaroufim @drisspg I have resolved the conficts. Look forward to see your new advice.

[drisspg]

Hey @HandH1998 rekicking off the internal CI/CD here that failed

Overall this works looks very good, will do a more thorough review

[drisspg]

It seems that you also applied formatting to a lot of files. This makes it pretty hard to review since all the changes get mixed together.

I opened #1226

Would you mind opening a separate PR for some of the other files you touched? And let's first add the formatting and then we can merge in your changes and make it easier to review.

[HandH1998]

@drisspg I have removed the extra formatting in this PR, which should now simplify the review process.

[alexsamardzic]

I've added a benchmarking script to #880, that makes it possible to compare the performance between the two W4A8 kernels. As CUTLASS-based version doesn't support group quantization, at the moment it is only possible to make the comparison with group_size=-1 in the Marlin-based version. Marlin-based version performs clearly better for input sizes less than 256, while CUTLASS-based version is faster for input sizes of 256 and higher. Consequently, Marlin-based version performs better on Llama generator too, with tokens/sec about 25% higher than CUTLASS-based version (note that in both cases I ran the generator as follows python generate.py --compile --precision torch.float16 -q ...). Please note that the comparison is not completely apples-to-apples, as besides group quantization support there are other small differences between kernels, but still it seems this is pretty much current state of affairs regarding the performance.

Let me return the compliment by stating that this Marlin-based kernel is a great work too. In particular, for me it clearly shows where the CUTLASS-based kernel should be improved.

[drisspg]

Overall looking really good, would you mind reporting lib size increase from this PR? I plan to take another once over of the cuda code tomorrow and then once CI is good this should be good to go :)

[HandH1998]

@drisspg @jerryzh168 Thanks for your reviews. I have resolved the most issues according to your advice. The increased lib size is about 5M. Let's go forward :)

[drisspg]

Lib increases from 3.5M -> 5.0M for provenance, I think this is acceptable

[drisspg]

@HandH1998 looks like there is 1 true failure on your PR would you mind fixing and then we can land?

[HandH1998]

@HandH1998 looks like there is 1 true failure on your PR would you mind fixing and then we can land?

I will try to fix it soon.

[jerryzh168]

can you add a table similar to https://github.com/pytorch/ao/tree/main/torchao/quantization#sparse-marlin to README to show the performance? otherwise looks good to me

[msaroufim]

just need to fix lint

[msaroufim]

note to self: this is the real code and it seems reasonable - rest is linting changes

[HandH1998]

can you add a table similar to https://github.com/pytorch/ao/tree/main/torchao/quantization#sparse-marlin to README to show the performance? otherwise looks good to me

I have added it.

[jerryzh168]

thanks!

[jerryzh168]

please fix lint: https://github.com/pytorch/ao/actions/runs/11812082822/job/32933511633?pr=1113

[zhyncs]

@HandH1998 It's so coooooool!