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.

❌ 11 New Failures

As of commit 8d68d45 with merge base 25ddb77 ():

NEW FAILURES - The following jobs have failed:

• 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.
• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t fb6dd0d6ae73ae1cbb60bac55e242a02d49789c07ca0a0d59e011095bfd8369a /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 1c7a6d90da8e32078590e8ae6a96776358576b868a95bfaf9e8ee537a7057a60 /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t f8361284dc6f056c297dfb1174c378a964e941c78ac609e09b6b1fe3922ef6e4 /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 657a2959f9c985513b47f64426869153ea5d9ffaccbac982359db62b1f410751 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 7da3b608584f4535272d57f663bbf750c3ae0255881315a8b39f63d91e5faeb1 /exec failed with exit code 1
• 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 3d785806e3d2fd810c77da83aa72456b12e66b6db6776fc8b2d8df72583494dd /exec failed with exit code 1
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch==2.7.0.dev20250122 --index-url https://down... / linux-job (gh)
  RuntimeError: Command docker exec -t 3d5c99040592bbb5a83d19e483dbc584f7db62f9518c32696f11465adec98a6b /exec failed with exit code 2
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch==2.7.0.dev20250122 --index-... / linux-job (gh)
  RuntimeError: Command docker exec -t 37e89103dea7c070525ab7a1aafcf412f3b99139b7edcb85ea5a821f86c0e6fc /exec failed with exit code 1
• Run TorchAO Experimental Tests / test (macos-14) (gh)
  Process completed with exit code 127.

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

[Degnel]

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.

Indeed, they are. It looks like only the CPU runs are failing. I presume that bitsandbytes might not install triton when no GPU is available (I might be missing something there). Here is an instance of a failing log:

https://github.com/pytorch/ao/actions/runs/13484452669/job/37730985419?pr=1763#step:14:1276

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

Thank you for the hint, I've locally updated the code accordingly 👍

[danielvegamyhre]

why is int8_dynamic_activation_int4_weight being used here?

[Degnel]

Thank's for noticing it. I was aiming for a static W4A8 quantization and I overlooked that it was dynamic. I will try to address this within the week.

[danielvegamyhre]

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.

Also @Degnel you should skip tests requiring triton if CUDA is not available.

[danielvegamyhre]

@Degnel thanks for your work on this, i ran the tests and it looks like your blockwise fp8 gemm test is failing due to quantization error