[New Feature] CUTLASS kernels for w4a8 quantization

[supriyar]

We plan to add QAT for LLMs to torchao (as mentioned in the original RFC here #47)

For this to run efficiently on the GPU we'd need kernel support for W4A8 quantization (int4 weights, int8 activations).

Other places where this has been raised before
NVIDIA/cutlass#1316,
NVIDIA/cutlass#1370

cc @andrewor14

[supriyar]

cc @alexsamardzic @cpuhrsch

[alexsamardzic]

Working on this: NVIDIA/cutlass#1413.

[jeromeku]

@alexsamardzic

Great work so far on integrating w4a8 GEMM in Cutlass!

Do you have plans on re-implementing this functionality in pre-Hopper architectures using Cutlass 3.x / CuTe rather the Cutlass 2.x apis that seem to be deprecated?

The 3.x interface has some convenient sub-byte primitives for slicing 4b tensors but warp-level shuffling would still be needed for efficient tensor core loading and mma.

Would be happy to help adapt 4b mixed type gemm using CuTe for Ampere.

[alexsamardzic]

Do you have plans on re-implementing this functionality in pre-Hopper architectures using Cutlass 3.x / CuTe rather the Cutlass 2.x apis that seem to be deprecated?

(Please send further comments to the PR mentioned above - I think it makes most sense to discuss CUTLASS features on CUTLASS GitHub pages.)

As it could be seen from my PR, this feature is implemented the same way as F16/S8, and alike. For my purpose, and that is adding support for this operation into PyTorch, for Ampere architecture and for both eager and compiled mode, this is good enough. I'm not sure in which way my changes could be made more 3.x-y, as the functionality is implemented on the warp level, but if you have any suggestions, please post them either into this, or in separate PR.