Add sparse marlin 2:4 gemm op

[Diogo-V]

Description

This PR is a more concise version of #621 where only the gemm op and surrounding functions are implemented for a 2:4 sparse marlin kernel.

What was done:

• Tests to validate the gemm op as well as with opcheck() to check if torch.compile will work out of the box with it
• Implemented functions to pack an int4 quantized tensor into a sparse marlin representation
• Implemented functions to reverse the above process (to be later used when dequantize() is called)

Notes:

• The cuda kernel was extracted from this repo

cc @jcaip

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 8699877 with merge base 0ed3090 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jcaip]

Thanks @Diogo-V!

[msaroufim]

that was fast haha, 1 shot green ci ;)

[Diogo-V]

Glad I could be of help!
Now, I have to maintain the streak ;)