Feat: Implementation of the DeepSeek blockwise quantization for fp8 tensors

[Degnel]

This PR is the first step towards addressing issue #1594. It includes the following implementations:

• fp8 triton gemm for blockwise quantisation
• quant, dequant and linear utilities
• time & precision benchmarks
• basic tests

If the code is validated, it would be great to bench it on H100.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[danielvegamyhre]

Thanks for your work on this! I'll take a closer look next week.

cc @vkuzo @drisspg

[Degnel]

Thanks for running the tests. I have two questions regarding the errors:

• Where should I add Triton to allow the tests to run successfully without introducing unnecessary dependencies in dev-requirements.txt?
• Does torchao provide any utility to check the available FP8 types for each gpu architecture?

[danielvegamyhre]

Thanks for running the tests. I have two questions regarding the errors:

• Where should I add Triton to allow the tests to run successfully without introducing unnecessary dependencies in dev-requirements.txt?

Can you clarify what you mean? Are tests failing in CI due to a missing triton installation? That shouldn't be happening, please share the link/logs if so.

• Does torchao provide any utility to check the available FP8 types for each gpu architecture?

We just use helpers which skip tests if GPU architecture is not at least SM 89:

ao/torchao/utils.py

Line 619 in f478692

def is_sm_at_least_89():

You can find examples in the float8 tests (example).