Fused triton quantize-dequantize in forward helpers

[ElizaWszola]

Create Triton kernel for dequantize op. Also make a kernel for fused quantize-dequantize.

While pure triton dequantize tends to be slower than the pytorch baseline, it is beneficial to implement a fused triton quantize-dequantize operation.

Unit tests

tests/test_quantization/lifecycle/test_forward.py -k test_dequantize_triton_matches_cpu
tests/test_quantization/lifecycle/test_forward.py -k test_quantize_dequantize_triton_matches_cpu

Benchmark output fused vs. unfused quantize-dequantize:

=================================================================================================================================================
SUMMARY - All times in ms, speedups relative to PyTorch baseline
=================================================================================================================================================
Config               Size         PyTorch    Mixed      T Unfused  T Fused    Mix SpUp  Unf SpUp  Fused SpUp  OK  
-------------------------------------------------------------------------------------------------------------------------------------------------
int8                 4096x4096       0.22 ms    0.16 ms    0.15 ms    0.10 ms   1.42x    1.50x    2.27x      Yes 
int8                 4096x11008      0.57 ms    0.33 ms    0.32 ms    0.18 ms   1.73x    1.78x    3.17x      Yes 
int8                 8192x8192       0.83 ms    0.46 ms    0.45 ms    0.24 ms   1.80x    1.86x    3.42x      Yes 
int4                 4096x4096       0.22 ms    0.15 ms    0.15 ms    0.10 ms   1.44x    1.47x    2.24x      Yes 
int4                 4096x11008      0.57 ms    0.33 ms    0.32 ms    0.18 ms   1.73x    1.76x    3.18x      Yes 
int4                 8192x8192       0.83 ms    0.46 ms    0.45 ms    0.24 ms   1.80x    1.85x    3.40x      Yes 
fp4                  4096x4096       0.22 ms    0.15 ms    0.15 ms    0.10 ms   1.45x    1.48x    2.21x      Yes 
fp4                  4096x11008      0.57 ms    0.33 ms    0.32 ms    0.19 ms   1.73x    1.76x    3.01x      Yes 
fp4                  8192x8192       0.83 ms    0.46 ms    0.45 ms    0.26 ms   1.79x    1.85x    3.27x      Yes 
fp8                  4096x4096       0.26 ms    0.15 ms    0.15 ms    0.10 ms   1.71x    1.76x    2.74x      Yes 
fp8                  4096x11008      0.67 ms    0.33 ms    0.32 ms    0.18 ms   2.01x    2.06x    3.64x      Yes 
fp8                  8192x8192       0.98 ms    0.46 ms    0.44 ms    0.24 ms   2.12x    2.20x    4.00x      Yes 
int8_channel         4096x4096       0.23 ms    0.14 ms    0.14 ms    0.08 ms   1.60x    1.61x    2.75x      Yes 
int8_channel         4096x11008      0.58 ms    0.32 ms    0.30 ms    0.17 ms   1.80x    1.90x    3.46x      Yes 
int8_channel         8192x8192       0.85 ms    0.46 ms    0.43 ms    0.23 ms   1.86x    1.97x    3.70x      Yes 
int4_channel         4096x4096       0.23 ms    0.14 ms    0.14 ms    0.08 ms   1.58x    1.64x    2.73x      Yes 
int4_channel         4096x11008      0.58 ms    0.32 ms    0.31 ms    0.17 ms   1.80x    1.88x    3.45x      Yes 
int4_channel         8192x8192       0.85 ms    0.46 ms    0.43 ms    0.23 ms   1.86x    1.97x    3.69x      Yes 
int8_g128            4096x4096       0.23 ms    0.19 ms    0.22 ms    0.12 ms   1.24x    1.03x    1.86x      Yes 
int8_g128            4096x11008      0.58 ms    0.35 ms    0.35 ms    0.20 ms   1.66x    1.65x    2.88x      Yes 
int8_g128            8192x8192       0.86 ms    0.60 ms    0.71 ms    0.37 ms   1.43x    1.20x    2.33x      Yes 
int4_g128            4096x4096       0.23 ms    0.18 ms    0.22 ms    0.12 ms   1.24x    1.02x    1.89x      Yes 
int4_g128            4096x11008      0.58 ms    0.35 ms    0.35 ms    0.20 ms   1.67x    1.67x    2.90x      Yes 
int4_g128            8192x8192       0.86 ms    0.60 ms    0.71 ms    0.37 ms   1.43x    1.20x    2.33x      Yes 
fp4_g32              4096x4096       0.23 ms    0.19 ms    0.22 ms    0.13 ms   1.24x    1.04x    1.80x      Yes 
fp4_g32              4096x11008      0.59 ms    0.37 ms    0.40 ms    0.23 ms   1.57x    1.48x    2.58x      Yes 
fp4_g32              8192x8192       0.86 ms    0.60 ms    0.72 ms    0.37 ms   1.42x    1.20x    2.30x      Yes 

=================================================================================================================================================
BENCHMARKED MODES
=================================================================================================================================================

  PyTorch:    PyTorch quantize + PyTorch dequantize (native PyTorch ops, baseline)
  Mixed:      Triton quantize + PyTorch dequantize
  T Unfused:  Triton quantize + Triton dequantize (two separate kernels)
  T Fused:    Triton quantize+dequantize (single fused kernel)

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Team

Run ID: bcdf0c0f-726f-4673-9676-762a72612250

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ElizaWszola.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yiliu30]

Hi @ElizaWszola, nice work! Would it make sense to add torch.compile as a baseline as well? It can automatically generate the Triton kernels and apply the necessary operator fusion.

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to adddress the lint failures. You will need to install the
dev optional install to get the required linting packages.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require one maintainer review	👀 reviews

🔴 Require one maintainer review

Waiting for any of

• approved-reviews-by=HDCharles
• approved-reviews-by=brian-dellabetta
• approved-reviews-by=dsikka
• approved-reviews-by=kylesayrs

This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

• any of:
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
• #changes-requested-reviews-by = 0

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ElizaWszola.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork