Check numerical equivalence / closeness between different kernel preferences

[jerryzh168]

Stacked PRs:

• Align Int4Tensor implementation details with the design of Float8Tensor #2687
• Add all fbgemm kernel Tensors into Int4WeightOnlyConfig and Float8DynamicActivationInt4WeightConfig #2474
• ->Check numerical equivalence / closeness between different kernel preferences #2651
• Bump version for float8 dynamic quant and weight only quant configs #2650
• Add Float8Tensor #2463

---

Check numerical equivalence / closeness between different kernel preferences

Summary:
This PR checks different kernel preferences for Float8Tensor are similar in numerics
(AUTO, TORCH and FBGEMM)

triton implementation and torchao implementation are a bit different right now actually, need to decide if we should fix it or not

1. difference in quantize op
   main difference seems to be the triton implementation is using:

a_scale = MAX_FP8 / max_abs
then do
a_scale = 1.0 / a_scale
a_fp8 = a * a_scale

while torch is doing:

a_scale = max_abs / MAX_FP8
a_fp8 = a / a_scale

Also the hp_value_lb and hp_value_ub settings are slightly different

triton choose scale and quantize code: https://github.com/pytorch/FBGEMM/blob/a4286c01ef01dad435b2ec8798605127d3032cd8/fbgemm_gpu/experimental/gemm/triton_gemm/fp8_gemm.py#L2382-L2392

torchao choose scale and quantize code:

ao/torchao/quantization/quant_primitives.py

Line 2183 in 3c466f8

def _choose_scale_float8(

ao/torchao/quantization/quant_primitives.py

Line 2283 in 3c466f8

def _quantize_affine_float8(

2. (potentially) difference in matrix multiplication ops

TORCH and AUTO/FBGEMM are using different quantized mm ops

Added a reverse option to bring sqnr closer:

granularity: PerTensor()  sizes: ((128,), 256, 128)  kp: KernelPreference.AUTO tensor(inf, device='cuda:0', dtype=torch.bfloat16)
granularity: PerTensor()  sizes: ((128,), 256, 128)  kp: KernelPreference.FBGEMM tensor(inf, device='cuda:0', dtype=torch.bfloat16)
.granularity: PerTensor()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.AUTO tensor(inf, device='cuda:0', dtype=torch.bfloat16)
granularity: PerTensor()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.FBGEMM tensor(inf, device='cuda:0', dtype=torch.bfloat16)
.granularity: PerRow()  sizes: ((128,), 256, 128)  kp: KernelPreference.AUTO tensor(inf, device='cuda:0', dtype=torch.bfloat16)
granularity: PerRow()  sizes: ((128,), 256, 128)  kp: KernelPreference.FBGEMM tensor(inf, device='cuda:0', dtype=torch.bfloat16)
.granularity: PerRow()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.AUTO tensor(64.5000, device='cuda:0', dtype=torch.bfloat16)
granularity: PerRow()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.FBGEMM tensor(68., device='cuda:0', dtype=torch.bfloat16)

Test Plan:
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_kernel_preference_numerical_equivalence

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ghstack-mergeability-check and Check labels failing with 'Resource not accessible by integration'

❌ 11 New Failures

As of commit bd0faa2 with merge base b757fb9 ():

NEW FAILURES - The following jobs have failed:

• Code Analysis with Ruff / build (3.9) (gh)
• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  RuntimeError: Command docker exec -t 322300278b330094670b35861751cc53dc8c117bd0a03b629924432bf74ba87b /exec failed with exit code 2
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t e8ba193ce484a9637c9823f8eac4b770c7d85ec4d0e7d3b43d213ed20a0b4d25 /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, cp... / linux-job (gh)
  RuntimeError: Command docker exec -t 68ed8170372443b0e518a176f3b4fe523ea85d9bf316b1a96e21a69b8b5eb2d6 /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 be3f8f48b17e29ed2b974269fc311b86b8117b1e6baa36ac9973fd41fedbc5fa /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 db6c35065c52e8574d3a834b1d8e1b1296e48a1ffe187f91a28b1d83727c3504 /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 5fc1f7d73dd2dc1e83235debac64cf3a1f97f856800f3f0c93a1d23ebb4f3c13 /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 50b29d6230e5717bf9fe51b16457adea943dc225b754630ca34594e63ec27c26 /exec failed with exit code 2
• 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 1721e7b360ac48e6daf91c6e678ae80ec3c7034bdda391b138f6fc9f67e8399c /exec failed with exit code 2
• 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 8a564c5037f92fb9d97dafae25a61b4ed1d4f55c39bf6d386b34393316f27f09 /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 abb6fa4a9f0919e3eebcab2c890c48d5ab01eecdb25c91fa602ce1509e9b68a5 /exec failed with exit code 2

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

[jerryzh168]

cc @vkuzo we don't have equivalence yet due to some differences in implementation, do you think we should match torchao quant primitives (choose_scale_float8 + quantize_float8) and triton ones?

[vkuzo]

do we know what the differences are?

IMO we should also choose either TORCH or FBGEMM (but not AUTO) as the reference, and match others to the reference

[jerryzh168]

yeah see PR summary for differences

I can update and use TORCH as reference

[vkuzo]

I think we should make the callsites match instead of having a flag

[jerryzh168]

yeah this is not the final state, was just trying out things to check if this helps or not