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 API	65.47 ms	234.39 ms
New API	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.

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 tensor slicing granularity. Without this change, we can't use per-tensor (0D scale) and per-row (1D scale).

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)

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

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

🔗 Helpful Links

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

Uh oh!

jerryzh168 Oct 24, 2025

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!

namgyu-youn Oct 24, 2025

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!

jerryzh168 Oct 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jerryzh168 left a comment

Choose a reason for hiding this comment

Uh oh!

jerryzh168 Oct 30, 2025

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!

jerryzh168 Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

jerryzh168 Oct 30, 2025

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

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

namgyu-youn 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