[WIP] 2:4 activation sparsity

[jcaip]

This PR is meant to give users the ability to accelerate LLMs with 2:4 activation sparsity, using the approach outlined in our ICLR workshop paper: https://arxiv.org/abs/2503.16672

The main contribution is a cutlass 24_fp8_pack kernel that is able to relatively efficiently calculate the packed data and metadata given a normal dense tensor.

The correctness of the kernel

Work that needs to be done:

[ ] selective weight loading kernel for fast decode?
[ ] Figure out some way to fuse the 24 sparsification into the fp8 rowwise quantization / activation function. We can either update the CUTLASS kernel or write a triton kernel to do the sparsification.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• [Infra] Jobs got intermittently cancelled/fail midway checkout

❌ 8 New Failures, 2 Unrelated Failures

As of commit 88bec35 with merge base 7513042 ():

NEW FAILURES - The following jobs have failed:

• Build Linux Wheels / build / build-manywheel-py3_9-cuda11_8 (gh)
• Build Linux Wheels / build / build-manywheel-py3_9-cuda12_6 (gh)
• Build Linux Wheels / build / build-manywheel-py3_9-cuda12_8 (gh)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cuda11_8 (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cuda12_6 (gh)
• Build Linux Wheels / build / upload / upload-manywheel-py3_9-cuda12_8 (gh)
• Code Analysis with Ruff / build (3.9) (gh)
  Process completed with exit code 1.
• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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