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.

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

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

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)

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

above comment is incorrect and this change is unrelated; #3241

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

this can be moved to torchao/quantization/quantize_/common/utils.py I think

Okay, then I will move this to torchao/quantization/quantize_/common/utils.py after this PR.

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

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.

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

namgyu-youn · 2025-11-11T10:20:29Z

Updated logs:

Set per-row as a default granularity in both api and test: introduce new int8 quantization API #3241 (comment) introduce new int8 quantization API #3241 (comment)
Update kernel (trition, int_mm) test using parameterization and FileCheck: introduce new int8 quantization API #3241 (comment) introduce new int8 quantization API #3241 (comment)
Fix quantization ops similar to UINT pattern: introduce new int8 quantization API #3241 (comment) introduce new int8 quantization API #3241 (comment) introduce new int8 quantization API #3241 (comment)

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

jerryzh168 · 2025-11-11T18:05:25Z

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

+            dtype=self.dtype,
+            device="cuda",
+        )
+        linear.weight.data = self.weight_fp.cuda()


why do we do this? doesn't sound very necessary?

jerryzh168 · 2025-11-11T18:06:27Z

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

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


nit: can you merge the checks in this test to previous test test_int8_linear_variants?

jerryzh168 · 2025-11-11T18:10:41Z

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

+            f"Dequantization error is too high to get a SQNR of {compute_error(dequantized, test_data)}"
+        )
+
+    @common_utils.parametrize(


does this have to be parametrize? I think what we need here is to check the code contains a sequence of ops / kernel calls, like this:

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

Lines 827 to 831 in b4ec4cb

FileCheck().check_count(

"torch.ops.triton.quantize_fp8_row.default(", 1

).check_count("torch.ops.fbgemm.f8f8bf16_rowwise.default(", 1).check_not(

".run("

).run(code[0])

I think we can check 1. the quantize op and then 2. the mm op extern_kernels._int_mm, in a single run (see example), that should be enough

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

jerryzh168 · 2025-11-11T18:48:42Z

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

+        self.weight_fp = torch.randn(*self.test_shape, dtype=self.dtype)
+        self.input_fp = torch.randn(*self.test_shape, dtype=self.dtype)
+        self.bias = torch.randn(self.test_shape[0], dtype=self.dtype)
+        self.block_size = list(self.test_shape)


I feel we probably don't need these, it's also easier for people to follow to define everything / most of things in the test itself

jerryzh168 · 2025-11-11T18:51:30Z

torchao/float8/inference.py


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


note: I think we can just check for ndim consistently everywhere, after #3324 is fixed

also isn't handling for per tensor and per row already included in original code?

ao/torchao/float8/inference.py

Lines 158 to 178 in b4ec4cb

if block_size_for_dim == 1:

# Scale is per-element along this dimension

# Slice away as normal

return aten.slice.Tensor(scale, dim, start, end, step)

else:

# There is blocking in this dimension

# Calculate which scale elements correspond to the sliced data

scale_start = start // block_size_for_dim if start is not None else None

scale_end = (

(end + block_size_for_dim - 1) // block_size_for_dim

if end is not None

else None

)

# Error on Step > 1

if step > 1:

raise NotImplementedError(

"Slicing with step > 1 is not implemented for scale tensors."

)

return aten.slice.Tensor(scale, dim, scale_start, scale_end, 1)

Reverted; I miscalculated block_sizes, and this change is unrelated. The original code is already handling per-tensor and per-row.

btw, this function might be improved by separating the granularity check I think; please check #3345 for this update.

jerryzh168 · 2025-11-11T18:56:24Z

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

+            if len(act_kwargs.block_size) != input_ndim:
+                if input_ndim == 3 and len(act_kwargs.block_size) == 2:
+                    block_size_updated = [1] + list(act_kwargs.block_size)
+                else:
+                    block_size_updated = list(act_kwargs.block_size)[-input_ndim:]


we kind of changed the meaning of block_size used in PerBlock quant recently, check

ao/torchao/quantization/utils.py

Line 707 in b4ec4cb

elif isinstance(granularity, PerBlock):

when is this code needed?

in principle we shouldn't update block_size here, but instead, make sure block_size make sense and is consistent throughout the code base

this comment is not addressed

jerryzh168 · 2025-11-11T18:58:30Z

torchao/quantization/quant_api.py

+        quantized_weight = Int8Tensor.from_hp(
+            weight,
+            block_size,
+            act_quant_kwargs=QuantizeTensorToInt8Kwargs(block_size=block_size),


why does activation use the same block_size as weight?

I think we can change block_size argument to granularity, if the block_size is unknown for activation since the shape is unknown

jerryzh168

please rebase, also some comments are not addressed yet I think

namgyu-youn · 2025-11-15T01:03:23Z

please rebase, also some comments are not addressed yet I think

Thanks for the fast review, but please wait until the review request. I will address all comments this weekend :)

jerryzh168 · 2025-11-18T23:22:26Z

torchao/quantization/quant_api.py

    """

    group_size: Optional[int] = None
+    granularity: Optional[Union[PerRow, PerTensor]] = PerRow()


nit: the type can just be Optional[Granularity] I think, we can do validation on supported granularity later

jerryzh168 · 2025-11-18T23:22:34Z

torchao/quantization/quant_api.py

    layout: Optional[Layout] = PlainLayout()
    act_mapping_type: Optional[MappingType] = MappingType.SYMMETRIC
    weight_only_decode: bool = False
+    granularity: Optional[Union[PerRow, PerTensor]] = PerRow()


same here

also should this one be a tuple? since there is both activation and weight

or follow this:

ao/torchao/quantization/quant_api.py

Line 1804 in 8d546c8

granularity: Optional[Union[FP8Granularity, List[FP8Granularity]]] = None

jerryzh168 · 2025-11-18T23:23:55Z

torchao/quantization/quant_api.py

+
+        assert config.version == 2, f"Unexpected version: {config.version}"
+        # Compute block_size from granularity for activation quantization kwargs
+        block_size = get_block_size(weight.shape, config.granularity)


this is incorrect I think, we don't want to calculate the block_size for activation from weight

jerryzh168 · 2025-11-18T23:24:14Z

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

+        # TODO: Static quantization support using `static_scale`, `static_zero_point`
+    """
+
+    block_size: list[int]


this should probably just take granularity

jerryzh168 · 2025-11-18T23:25:17Z

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

+        while scale.ndim < qdata_fp.ndim:
+            scale = scale.unsqueeze(-1)
+
+        scale_expanded = _maybe_expand_scale_to_tensor_shape(scale, qdata_fp.shape)


are these needed?

jerryzh168 · 2025-11-18T23:25:57Z

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

+            scale = scale.unsqueeze(-1)
+
+        scale_expanded = _maybe_expand_scale_to_tensor_shape(scale, qdata_fp.shape)
+        return qdata_fp * scale_expanded.to(output_dtype)


I think we should just call the dequantize_affine op directly

ao/torchao/quantization/quant_primitives.py

Line 704 in 8d546c8

def dequantize_affine(

jerryzh168 · 2025-11-18T23:26:51Z

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

+        # can downcast only at the very end
+        output_dtype = activation_tensor.dtype
+        y = y.to(output_dtype)


nit: can probably do this after bias addition

jerryzh168 · 2025-11-18T23:27:02Z

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

+        y = y.reshape(*activation_tensor.shape[:-1], weight_tensor.qdata.shape[0])
+        if bias is not None:
+            y += bias
+        return y


can do the cast right before return I think

jerryzh168

thanks for the updates, test looks good I think, still some comments that need to be addressed in int8_tensor.py

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

test/quantization/quantize_/workflows/int8/test_int8_tensor.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Oct 24, 2025

View reviewed changes

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

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

test/quantization/quantize_/workflows/int8/test_int8_tensor.py Outdated Show resolved Hide resolved

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

test/quantization/quantize_/workflows/int8/test_int8_tensor.py Outdated Show resolved Hide resolved

fix quantization ops

3ab38ba

namgyu-youn requested a review from jerryzh168 November 11, 2025 10:20

jerryzh168 reviewed Nov 11, 2025

View reviewed changes

test/quantization/quantize_/workflows/int8/test_int8_tensor.py Show resolved Hide resolved

jerryzh168 reviewed Nov 11, 2025

View reviewed changes

test/quantization/quantize_/workflows/int8/test_int8_tensor.py Show resolved Hide resolved

jerryzh168 reviewed Nov 11, 2025

View reviewed changes

merge test cases with cleanup

b516304

jerryzh168 requested changes Nov 14, 2025

View reviewed changes

update block_size args to granularity

0c2bb76

namgyu-youn force-pushed the int8-quant-api branch from 425f554 to 0c2bb76 Compare November 16, 2025 13:28

update expected kernel test

d11af10

namgyu-youn force-pushed the int8-quant-api branch from 42cdb05 to d11af10 Compare November 16, 2025 14:09

Merge branch 'main' into int8-quant-api

027afd8

namgyu-youn requested a review from jerryzh168 November 16, 2025 14:23

jerryzh168 reviewed Nov 18, 2025

View reviewed changes

jerryzh168 requested changes Nov 18, 2025

View reviewed changes

	FileCheck().check_count(
	"torch.ops.triton.quantize_fp8_row.default(", 1
	).check_count("torch.ops.fbgemm.f8f8bf16_rowwise.default(", 1).check_not(
	".run("
	).run(code[0])

	if block_size_for_dim == 1:
	# Scale is per-element along this dimension
	# Slice away as normal
	return aten.slice.Tensor(scale, dim, start, end, step)
	else:
	# There is blocking in this dimension
	# Calculate which scale elements correspond to the sliced data
	scale_start = start // block_size_for_dim if start is not None else None
	scale_end = (
	(end + block_size_for_dim - 1) // block_size_for_dim
	if end is not None
	else None
	)

	# Error on Step > 1
	if step > 1:
	raise NotImplementedError(
	"Slicing with step > 1 is not implemented for scale tensors."
	)

	return aten.slice.Tensor(scale, dim, scale_start, scale_end, 1)

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

Uh oh!

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!

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!

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!

Uh oh!

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

Uh oh!

Uh oh!

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!

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

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

namgyu-youn commented Nov 11, 2025 •

edited

Loading

namgyu-youn Nov 16, 2025 •

edited

Loading

jerryzh168 Nov 18, 2025 •

edited

Loading

jerryzh168 left a comment •

edited

Loading