Add selective weight loading decode kernel for activation sparsity

[jcaip]

Summary:

This PR adds in a kernel to accelerate $$xW^T$$ when x is sparse and we are memory bound.

The idea here is that we can avoid loading the columns of $W$ that correspond to the zero elements of $x$.

This lets us accelerate activation sparsity for bs=1 decode use cases. The kernel is an adapted version of https://github.com/FasterDecoding/TEAL

Test Plan:
pytest test/sparsity/test_activation24.py

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 9 New Failures

As of commit 9e33323 with merge base 4c6188f ():

NEW FAILURES - The following jobs have failed:

• Code Analysis with Ruff / build (3.9) (gh)
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cp... / linux-job (gh)
  RuntimeError: Command docker exec -t 845311ca209794f542a50be47c5c62b84d027aca2ebd97c178460b98cc85ad5e /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.6, linux.4xlarge, torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 58699916c745315b380667ad733ed47c9cbd2ea3ca81153e72769aa67a698abe /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t aa21594709c3ced3b0b3a4e4fa7068fb38f6b1e825bd0f1999e06e0729998d44 /exec failed with exit code 2
• 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)
  RuntimeError: Command docker exec -t e0821c015bf56dd992fc7376c0664669bda9464522c4e19600509eb267b400aa /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 17ca13e332bffd849fd6624b7c9495b7b2fa6f5e32b984906d66da7661d73e5c /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 1050a9d5e0b4c5e3c93cb189cadb3e75bd951cf7aa3167921fa45931863e849c /exec failed with exit code 1
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t 6a5560aa41e6568164fdc400ec85878bc1daaad2408b2486e61e6edc7998e6be /exec failed with exit code 2
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t b0c99151c471bf72c20e7b56816712b600bd26e4301877e043ac70806b52eca1 /exec failed with exit code 1

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