2:4 activation sparsity packing kernels

[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, which I've copied over from xFormers.

Performance Benchmarks

python benchmarks/benchmark_e2e_fp8_sparse_linear.py 

|   num_tokens |   bf16_latency (us) |   bf16_c_latency (us) |   fp8_c_time (us) |   fp8_c_sparse_time (us) |   fp8_c_activation_sparse_time (us) |   speedup |
|-------------:|--------------------:|----------------------:|------------------:|-------------------------:|------------------------------------:|----------:|
|           64 |             166.816 |               163.04  |           103.008 |                   74.304 |                             102.816 |  1.00187  |
|          128 |             156.256 |               151.52  |            99.936 |                   75.456 |                             102.048 |  0.979304 |
|          256 |             172.288 |               159.584 |           114.08  |                   82.432 |                             111.072 |  1.02708  |
|          512 |             218.88  |               204.608 |           144.096 |                  114.56  |                             139.488 |  1.03304  |
|         1024 |             394.4   |               392.544 |           251.104 |                  196.416 |                             227.904 |  1.1018   |
|         2048 |             764.608 |               734.816 |           480.704 |                  381.152 |                             426.688 |  1.12659  |
|         4096 |            1658.82  |              1623.58  |           901.344 |                  779.008 |                             843.392 |  1.06871  |

Tests

pytest tests/sparsity/test_activation24.py

[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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 126166f with merge base 66eb801 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh) (trunk failure)
  Process completed with exit code 1.

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