Check numerical equivalence / closeness between different kernel preferences #2651

jerryzh168 · 2025-08-01T00:52:56Z

Stacked PRs:

Support optional_tensor_names in TorchAOBaseTensor #2710
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

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

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(

(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 · 2025-08-01T00:53:00Z

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

⏳ No Failures, 1 Pending

As of commit 1506c0d with merge base d2e791b ():
💚 Looks good so far! There are no failures yet. 💚

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

jerryzh168 · 2025-08-01T00:54:43Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

+        for i in range(1, len(kp_and_res)):
+            kp, res = kp_and_res[i]
+            self.assertTrue(
+                compute_error(res, kp_and_res[0][1]) > 28,


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?

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

yeah see PR summary for differences

I can update and use TORCH as reference

torchao/quantization/quant_primitives.py

…erences 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: https://github.com/pytorch/ao/blob/3c466f844684af0fb80014094f2ca8663881eb33/torchao/quantization/quant_primitives.py#L2183 https://github.com/pytorch/ao/blob/3c466f844684af0fb80014094f2ca8663881eb33/torchao/quantization/quant_primitives.py#L2283 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: stack-info: PR: #2651, branch: jerryzh168/stack/15

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

drisspg · 2025-08-07T03:26:16Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

+
+        # comparing numerics between different kernel preferences, using TORCH as the standard
+        kp_and_res = list(quantized_outputs.items())
+        for i in range(1, len(kp_and_res)):


can you explicitly peel the first iteration so its very obvious that we see its the reference and just iterate on the rest of the keys

OK updated to run the ref separately

…erences 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: https://github.com/pytorch/ao/blob/3c466f844684af0fb80014094f2ca8663881eb33/torchao/quantization/quant_primitives.py#L2183 https://github.com/pytorch/ao/blob/3c466f844684af0fb80014094f2ca8663881eb33/torchao/quantization/quant_primitives.py#L2283 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: stack-info: PR: #2651, branch: jerryzh168/stack/15

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 2534529 to c608b78 Compare August 1, 2025 00:53

jerryzh168 force-pushed the jerryzh168/stack/14 branch from e19cb46 to 5ae457c Compare August 1, 2025 00:53

This was referenced Aug 1, 2025

Add Float8Tensor #2463

Merged

Bump version for float8 dynamic quant and weight only quant configs #2650

Merged

Add all fbgemm kernel Tensors into Int4WeightOnlyConfig and Float8DynamicActivationInt4WeightConfig #2474

Merged

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 1, 2025

jerryzh168 added the topic: not user facing Use this tag if you don't want this PR to show up in release notes label Aug 1, 2025

jerryzh168 commented Aug 1, 2025

View reviewed changes

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 1, 2025 00:56

jerryzh168 force-pushed the jerryzh168/stack/15 branch from c608b78 to 42a767c Compare August 1, 2025 00:56

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 1, 2025 00:56

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 1, 2025 03:38

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 42a767c to 65a4f84 Compare August 1, 2025 03:38

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 1, 2025 03:38

jerryzh168 requested a review from vkuzo August 1, 2025 04:43

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 1, 2025 21:12

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 65a4f84 to ba8efe2 Compare August 1, 2025 21:13

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 1, 2025 21:13

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 2, 2025 01:31

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 2, 2025 01:31

vkuzo reviewed Aug 4, 2025

View reviewed changes

torchao/quantization/quant_primitives.py Outdated Show resolved Hide resolved

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 4, 2025 17:30

jerryzh168 force-pushed the jerryzh168/stack/15 branch from ba8efe2 to 1720743 Compare August 4, 2025 17:30

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 4, 2025 17:30

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 4, 2025 18:15

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 1720743 to 36fce5e Compare August 4, 2025 18:15

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 4, 2025 18:15

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 4, 2025 22:14

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 36fce5e to 9824504 Compare August 4, 2025 22:14

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 60eae40 to 2e63f70 Compare August 6, 2025 19:12

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 6, 2025 19:12

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 6, 2025 21:18

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 2e63f70 to 48bae28 Compare August 6, 2025 21:18

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 6, 2025 21:18

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 6, 2025 22:09

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 48bae28 to ab6d944 Compare August 6, 2025 22:10

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 6, 2025 22:10

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 6, 2025 22:16

jerryzh168 force-pushed the jerryzh168/stack/15 branch from ab6d944 to 57b2086 Compare August 6, 2025 22:16

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 6, 2025 22:16

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 6, 2025 23:27

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 57b2086 to f6bac49 Compare August 6, 2025 23:27

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 6, 2025 23:27

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 7, 2025 02:57

jerryzh168 force-pushed the jerryzh168/stack/15 branch from f6bac49 to c7f8ff0 Compare August 7, 2025 02:57

jerryzh168 changed the base branch from main to jerryzh168/stack/14 August 7, 2025 02:57

jerryzh168 mentioned this pull request Aug 7, 2025

Support optional_tensor_names in TorchAOBaseTensor #2710

Merged

jerryzh168 force-pushed the jerryzh168/stack/15 branch from c7f8ff0 to 847259b Compare August 7, 2025 02:58

jerryzh168 changed the base branch from jerryzh168/stack/14 to main August 7, 2025 02:58

drisspg reviewed Aug 7, 2025

View reviewed changes

test/quantization/quantize_/workflows/float8/test_float8_tensor.py Outdated Show resolved Hide resolved

drisspg reviewed Aug 7, 2025

View reviewed changes

jerryzh168 force-pushed the jerryzh168/stack/15 branch 3 times, most recently from 9f5fb3f to 873b120 Compare August 7, 2025 04:29

jerryzh168 force-pushed the jerryzh168/stack/15 branch from 873b120 to 1506c0d Compare August 7, 2025 20:56

andrewor14 approved these changes Aug 7, 2025

View reviewed changes

jerryzh168 merged commit 1114ca0 into main Aug 7, 2025
19 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Check numerical equivalence / closeness between different kernel preferences #2651

Check numerical equivalence / closeness between different kernel preferences #2651

Uh oh!

jerryzh168 commented Aug 1, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Aug 1, 2025 •

edited

Loading

Uh oh!

jerryzh168 Aug 1, 2025

Uh oh!

vkuzo Aug 1, 2025

Uh oh!

jerryzh168 Aug 1, 2025

Uh oh!

Uh oh!

Uh oh!

drisspg Aug 7, 2025

Uh oh!

jerryzh168 Aug 7, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Check numerical equivalence / closeness between different kernel preferences #2651

Check numerical equivalence / closeness between different kernel preferences #2651

Uh oh!

Conversation

jerryzh168 commented Aug 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!