introduce new int8 quantization API #3241

namgyu-youn · 2025-10-24T12:50:23Z

Summary:
Introduce a new tensor subclass API. Main features are

Int8Tensor: Main API, which handles quantization and dequantization operations
Utility operation functions: Tensor slice, index selection

This api is integrated to global variants (Int8WeightOnlyConfig, Int8DynamicActivationInt8WeightConfig) using version, and not defined as a default.

Related Issue/PR:
This is reopened PR for #3038

Test plan:
test/quantization/quantize_/workflows/int8/test_int8_tensor.py

Performance:
The following are the results of https://github.com/pytorch/ao/blob/main/tutorials/quantize_vit/run_vit_b_quant.py with a batch size of 32:

API	With `torch.compile`	Without `torch.compile`
Old	65.47 ms	234.39 ms
New	63.30 ms	239.30 ms

pytorch-bot · 2025-10-24T12:50:27Z

🔗 Helpful Links

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

📄 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:

ROCm failures during provisioning step due to network issues

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

jerryzh168 · 2025-10-24T17:39:48Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+            quant_min=self.int8_min,
+            quant_max=self.int8_max,


nit: we can omit these two args if these are the same as default (-128, 127)

jerryzh168 · 2025-10-24T17:43:01Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+        )
+
+    @common_utils.parametrize("dtype", [torch.bfloat16, torch.float16])
+    def test_quantization_shapes(self, dtype):


this seems to be a combination of two tests, one for dynamic quant one for static quant, can you use something like this:

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

Line 65 in e9c7bea

@common_utils.parametrize("mode", ["dynamic", "weight-only"])

also I feel it might be better to not add static quant in this PR, and in a separate PR add both the tensor support and config support for static quant

Okay, not sure to remove static flags (although its not fully implemented) before, but small PR should be always better I feel. I will remove static_scale and all those supports.

jerryzh168 · 2025-10-24T17:44:32Z

torchao/quantization/quantize_/workflows/int8/int8_tensor.py

+        if act_quant_kwargs is not None and act_quant_kwargs.static_scale is not None:
+            # INT8 × INT8 (static)
+            scale = act_quant_kwargs.static_scale
+            zero_point = torch.zeros_like(scale, dtype=torch.int8)


I think user should specify static_zero_point as well

but again, it's better to do this in a separate PR, since current state is a half of the static quant feature (no config)

jerryzh168 · 2025-10-24T17:46:37Z

torchao/quantization/quantize_/workflows/int8/int8_tensor.py

+        # Cast fp16 scale to float
+        intermediate_dtype = (
+            torch.float if x_scales.dtype == torch.half else x_scales.dtype
+        )
+        # Note: CUDA doesn't support int32/int64 matmul, so we convert to float
+        # Error message is NotImplementedError: "addmm_cuda" not implemented for 'Int'
+        # This may introduce minor numerical differences compared to int arithmetic
+        y_dot = torch.mm(tmp.to(intermediate_dtype), w_vals_t.to(intermediate_dtype))
+
+        # Apply activation scale
+        is_per_tensor_act = x_scales.numel() == 1
+        if is_per_tensor_act:
+            y_dot.mul_(x_scales.to(intermediate_dtype))
+        else:
+            # For block-wise activation scale, reshape to match y_dot
+            x_scales_reshaped = x_scales.view(y_dot.shape[0], -1)
+            y_dot.mul_(x_scales_reshaped.to(intermediate_dtype))
+
+        # Apply weight scale
+        is_per_tensor_weight = w_scales.numel() == 1
+        if is_per_tensor_weight:
+            result = y_dot.mul_(w_scales.to(intermediate_dtype))
+        else:
+            # Per-row weight scale - transpose and broadcast
+            w_scales_broadcast = w_scales.t().expand_as(y_dot)
+            result = y_dot.mul_(w_scales_broadcast.to(intermediate_dtype))
+
+        # Reshape back to original shape
+        result = result.view(*x_vals.shape[:-1], result.shape[-1])
+        result = result.to(activation_tensor.dtype)


this should follow:

ao/torchao/dtypes/uintx/plain_layout.py

Line 281 in e9c7bea

def _linear_int8_act_int8_weight_impl(input_tensor, weight_tensor, bias):

torchao/quantization/quantize_/workflows/int8/int8_tensor.py

jerryzh168

I think we should

split the static quant support to separate PR
follow what https://github.com/pytorch/ao/blob/main/torchao/dtypes/uintx/plain_layout.py is doing for quantized linear implementation

this should be a refactor PR, not a refactor + some extra modifications + some feature implementations I think

jerryzh168 · 2025-10-30T22:45:53Z

torchao/float8/inference.py

    aten = torch.ops.aten

-    # Unsupported case for now, this would be 1 scale per data element
+    # Per-tensor quantization (scalar scale)


is this change related?

It is updated to support more granularity. Without this change, we can't use per-tensor (0D scale) and per-row (1D scale).

So maybe it's better to move this util function to a common place?

jerryzh168 · 2025-10-30T22:51:04Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+            Int8WeightOnlyConfig(version=2),
+        ],
+    )
+    def test_per_row_scale_shape(self, dtype, config):


can you add a test like this one

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

Line 80 in 1e473ed

def test_fp8_linear_variants(

to test all the variations and include the shape check there?

jerryzh168 · 2025-10-30T22:51:47Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+
+    @common_utils.parametrize("dtype", [torch.bfloat16, torch.float16])
+    @common_utils.parametrize("has_bias", [True, False])
+    def test_weight_only_linear_with_bias(self, dtype, has_bias):


this can probably be merged into the linear varaints test as well

jerryzh168

thanks, I think the tensor changes looks good, but need to make a linear_variants tests to make sure we cover different aspects of things (e.g. compile), see comments inline

can you also do a e2e perf check with https://github.com/pytorch/ao/blob/main/tutorials/quantize_vit/run_vit_b_quant.py to make sure the performance are the same before and after change for vit model?

also add a kernel check might be useful to make sure we don't regress things:

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

Line 422 in 1e473ed

def test_expected_gpu_kernel_fbgemm(self):

namgyu-youn · 2025-10-31T10:06:08Z

Updated logs:

Support 3D Tensor input using reshape, similar to Flot8Tensor
Build integrated linear variants test: introduce new int8 quantization API #3241 (comment), introduce new int8 quantization API #3241 (comment)
Add available kernel test: introduce new int8 quantization API #3241 (review)
Add PERF result in PR contents: introduce new int8 quantization API #3241 (review)

Xia-Weiwen · 2025-11-03T05:43:10Z

Hi @namgyu-youn Do you plan to submit another PR for static quantization? We also need static quantization for SmoothQuant. So, we are wondering if you have a plan or we should consider adding it ourselves. Thanks. CC @cyxlily

namgyu-youn · 2025-11-03T06:37:52Z

Hi @namgyu-youn Do you plan to submit another PR for static quantization? We also need static quantization for SmoothQuant. So, we are wondering if you have a plan or we should consider adding it ourselves. Thanks. CC @cyxlily

Yeah, static quantization support using static/dynamic flags is planned; I hope to show it to your team in the foreseeable future.

Also, in the SmoothQuant case, validating its support for the new quantization APIs (below) has higher priority, I think. Could you look into it?

W4A16-INT: Int4WeightOnlyConfig(group_size=32, version=2)
W4A16-FP: Float8WeightOnlyConfig(version=2)
W8A8-FP-dynamic: Float8DynamicActivationFloat8WeightConfig(version=2)

Xia-Weiwen · 2025-11-03T06:52:03Z

Yeah, static quantization support using static/dynamic flags is planned; I hope to show it to your team in the foreseeable future.

Thanks. Looking forward to it. If there is anything we can help with, please let us know.

Also, in the SmoothQuant case, validating its support for the new quantization APIs (below) has higher priority, I think. Could you look into it?

W4A16-INT: Int4WeightOnlyConfig(group_size=32, version=2)

W4A16-FP: Float8WeightOnlyConfig(version=2)

W8A8-FP-dynamic: Float8DynamicActivationFloat8WeightConfig(version=2)

By "validating them", do you mean adding test cases? And are W4A16 and W8A16 (I guess there is a typo in your comment) really needed for SmoothQuant? For W4A16 , it would be much the same as AWQ. And for W8A16, I think accuracy is generally good enough without SmoothQuant.

namgyu-youn · 2025-11-03T07:36:21Z

By "validating them", do you mean adding test cases? And are W4A16 and W8A16 (I guess there is a typo in your comment) really needed for SmoothQuant? For W4A16 , it would be much the same as AWQ. And for W8A16, I think accuracy is generally good enough without SmoothQuant.

Oh yes, it was a typo (W8A16 is right), and W4A16-INT (Int4WeightOnlyConfig(group_size=32, version=2)) is of interest. In my last experience and https://arxiv.org/html/2411.02355v3, W4A16-INT is the most efficient choice for synchronous deployments, while W8A8-INT maximize throughput in asynchronous settings.

Because current AWQ/SmoothQuant test is only working with old APIs (version 1), we can replace it with new APIs like Int4WeightOnlyConfig(group_size=32, version=2) I guess.

Xia-Weiwen · 2025-11-03T07:59:00Z

By "validating them", do you mean adding test cases? And are W4A16 and W8A16 (I guess there is a typo in your comment) really needed for SmoothQuant? For W4A16 , it would be much the same as AWQ. And for W8A16, I think accuracy is generally good enough without SmoothQuant.

Oh yes, it was a typo (W8A16 is right), and W4A16-INT (Int4WeightOnlyConfig(group_size=32, version=2)) is of interest. In my last experience and https://arxiv.org/html/2411.02355v3, W4A16-INT is the most efficient choice for synchronous deployments, while W8A8-INT maximize throughput in asynchronous settings.

Because current AWQ/SmoothQuant test is only working with old APIs (version 1), we can replace it with new APIs like Int4WeightOnlyConfig(group_size=32, version=2) I guess.

I see. Thanks. We will evaluate that.

Xia-Weiwen · 2025-11-04T01:32:58Z

Hi @namgyu-youn May I know if you have a timeline to land this? Thanks.

jerryzh168 · 2025-11-04T01:41:59Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+            ((32, 128), 64, 256),  # 3D
+        ],
+    )
+    def test_int8_linear_quantization_accuracy(


can this one be combined with test_int8_linear_variants as well?

jerryzh168 · 2025-11-04T01:42:46Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+        with self.assertRaises(NotImplementedError):
+            _ = dummy.weight[::2]
+
+    def test_index_select(self):


I think all the tests should be modified to use config as well

jerryzh168 · 2025-11-04T01:44:08Z

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

+        kernels = {}
+
+        # Check for Triton kernels
+        if "torch.ops.triton" in code[0]:


should add some asserts I think

- Configs are updated to global variants

namgyu-youn · 2025-11-04T14:05:04Z

Updated logs:

Replaced Int8Tensor to configs in unit test: introduce new int8 quantization API #3241 (comment)
Updated linear variation tests: introduce new int8 quantization API #3241 (comment)
Update kernel detection test: introduce new int8 quantization API #3241 (comment)

introduce new int8 quantization API

da7cfea

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 Oct 24, 2025

namgyu-youn mentioned this pull request Oct 24, 2025

Add Int8Tensor for clearer interface #3038

Closed

jerryzh168 reviewed Oct 24, 2025

View reviewed changes

torchao/quantization/quantize_/workflows/int8/int8_tensor.py Show resolved Hide resolved

jerryzh168 requested changes Oct 24, 2025

View reviewed changes

namgyu-youn added 2 commits October 26, 2025 14:59

refactor ops and update test cases

27076d3

update granularity slicing support

cdb1d9f

namgyu-youn requested a review from jerryzh168 October 26, 2025 07:29

jerryzh168 reviewed Oct 30, 2025

View reviewed changes

namgyu-youn added 2 commits October 31, 2025 18:29

add 3D support to api, build linear variants test

9f1b6c9

add kernel detection test case

9301717

namgyu-youn requested a review from jerryzh168 October 31, 2025 09:57

refactor kernel test

caaba7a

jerryzh168 mentioned this pull request Oct 31, 2025

Migrating from AffineQuantizedTensor + Layouts to new structure of tensor subclasses #2752

Open

17 tasks

jerryzh168 reviewed Nov 4, 2025

View reviewed changes

update linear variant, kernel detection test

0f51ee6

- Configs are updated to global variants

namgyu-youn requested a review from jerryzh168 November 4, 2025 14:05

introduce new int8 quantization API #3241

Are you sure you want to change the base?

introduce new int8 quantization API #3241

Conversation

namgyu-youn commented Oct 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jerryzh168 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

namgyu-youn Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jerryzh168 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

namgyu-youn commented Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Xia-Weiwen commented Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

namgyu-youn commented Nov 3, 2025

Uh oh!

Xia-Weiwen commented Nov 3, 2025

Uh oh!

namgyu-youn commented Nov 3, 2025

Uh oh!

Xia-Weiwen commented Nov 3, 2025

Uh oh!

Xia-Weiwen commented Nov 4, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

namgyu-youn commented Nov 4, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

namgyu-youn commented Oct 24, 2025 •

edited

Loading

pytorch-bot bot commented Oct 24, 2025 •

edited

Loading

jerryzh168 Oct 24, 2025 •

edited

Loading

jerryzh168 Oct 24, 2025 •

edited

Loading

namgyu-youn Oct 31, 2025 •

edited

Loading

jerryzh168 left a comment •

edited

Loading

namgyu-youn commented Oct 31, 2025 •

edited

Loading

Xia-Weiwen commented Nov 3, 2025 •

edited

Loading