Add NPU (Ascend) backend support for INT4 weight-only quantization workflow #3172

orangeH25 · 2025-10-14T11:44:22Z

Related to #3044

Summary

This PR adds NPU (Ascend) backend support for the INT4 weight-only quantization workflow.

It introduces a new tensor subclass, Int4PlainInt32TensorNPU, aligned with the existing Int4PlainInt32Tensor for the plain_int32 packing format.

Environment

torchao version: 0.13.0 (main branch, commit: f64daac)
torch version: 2.7.1
torch_npu version: 2.7.1rc1
Ascend Toolkit (CANN): 8.2.RC1
Device: Ascend 910B4
OS: EulerOS 2.10 (Kernel 4.19.90, aarch64)
Python: 3.11

Files changed

Modified

torchao/quantization/__init__.py
torchao/quantization/quant_api.py
torchao/quantization/quantize_/workflows/__init__.py

Added

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor_npu.py
test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor_npu.py

Implementation Overview

Introduces Int4PlainInt32TensorNPU to enable NPU backend support for INT4 weight-only quantization.
Registeres new tensor subclass and integrated into quant_api.py for dispatch.
Updates __init__.py files to ensure proper import and exposure.
Adds corresponding test cases for NPU workflow.

Test Case

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor_npu.py

…rkflow

pytorch-bot · 2025-10-14T11:44:26Z

🔗 Helpful Links

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

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

fffrog · 2025-10-14T13:47:30Z

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor_npu.py

+try:
+    import torch_npu
+except ImportError:
+    torch_npu = None
+


PyTorch provide Autoload mechinasm, so we do not need to import it explicitly.

fffrog · 2025-10-14T13:50:58Z

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor_npu.py

+@unittest.skipIf(torch_npu is None, "torch_npu is not available")
+@unittest.skipIf(not torch_npu.npu.is_available(), "NPU not available")


Suggested change

@unittest.skipIf(torch_npu is None, "torch_npu is not available")

@unittest.skipIf(not torch_npu.npu.is_available(), "NPU not available")

@unittest.skipIf(torch.accelerator.current_accelerator(True).type == "npu" and torch.accelerator.is_available(), "NPU not available")

fffrog · 2025-10-14T13:56:02Z

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor_npu.py

+@unittest.skipIf(
+    version.parse(torch_npu.__version__) < version.parse("2.7.1rc1"),
+    "Need torch_npu 2.7.1rc1+",
+)


We can remove it because there are some strcit version mapping between PyTorch and Torch_NPU

fffrog · 2025-10-14T13:57:07Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor_npu.py

+        )
+
+        assert int_data.dtype == torch.int32, (
+            f"torch_npu.npu_convert_weight_to_int4pack expects `int32` dtype"


Suggested change

f"torch_npu.npu_convert_weight_to_int4pack expects `int32` dtype"

f"torch.ops.npu.npu_convert_weight_to_int4pack expects `int32` dtype"

fffrog · 2025-10-14T13:57:22Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor_npu.py

+        )
+
+        assert int_data.shape[-1] % 8 == 0, (
+            f"torch_npu.npu_convert_weight_to_int4pack expects last dim must be aligned to 8,but got {int_data.shape[-1]}"


Suggested change

f"torch_npu.npu_convert_weight_to_int4pack expects last dim must be aligned to 8,but got {int_data.shape[-1]}"

f"torch.ops.npu.npu_convert_weight_to_int4pack expects last dim must be aligned to 8,but got {int_data.shape[-1]}"

orangeH25 · 2025-10-15T06:30:26Z

Hi @jcaip @jerryzh168 , please help to review it, thanks!

jcaip · 2025-10-20T14:32:30Z

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor_npu.py

+    and torch.accelerator.is_available(),
+    "NPU not available",
+)
+class Int4PlainInt32TensorNPU(TestCase):


Just curious, do we need NPUs to test this? I don't think we have any in CI.

jcaip

Thanks for the PR @orangeH25 @fffrog!

The code looks good to me, but I'm curious on how to best test this? It looks like we skip tests in CI because we don't have NPU devices. I believe that NPU support was added to TorchTune as well, do you know how they test device specific functionality there?

Also, just a heads up most of the team is at PTC / Open source AI week in SF this week, so we might be a little slow in responding :)

jerryzh168 · 2025-10-20T21:42:15Z

please don't include device NPU in the name of Tensor, since packing_format is supposed to be agnostic to device

jerryzh168 · 2025-10-20T21:44:23Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor_npu.py

+    int4 weight-only quantization on Ascend NPU backend (groupwise quantization only)
+
+    Tensor Attributes:
+        qdata: (N, K/8), packed int4 weight, the data type is int32 here with 8*int4, the original dtype can be float16 or bfloat16


does this exactly align with Int4PlainInt32Tensor? if so, please merge with that tensor subclass

orangeH25 · 2025-10-21T09:12:12Z

Hi @jcaip @jerryzh168 ,thanks for the review!

Just curious, do we need NPUs to test this? I don't think we have any in CI.

Yes, this case is actually pretty common in open-source projects.

A typical approach is to set up a nightly CI job in the downstream repo that automatically pulls the latest code from the upstream repo each day for build and testing. The results are then shown as a GitHub badge in the upstream repo’s README.md to clearly show the latest build status.

does this exactly align with Int4PlainInt32Tensor? if so, please merge with that tensor subclass

You mean that we should keep the entry logic in quant_api.py unchanged:

elif int4_packing_format == Int4PackingFormat.PLAIN_INT32:
    new_weight = Int4PlainInt32Tensor.from_hp(
        weight,
        block_size,
    )
    return new_weight

and then handle different backend implementations in the from_hp and linear methods of Int4PlainInt32Tensor using simple if/else branches.

class Int4PlainInt32Tensor(TorchAOBaseTensor):
    ...
    @classmethod
    def from_hp(
        cls,
        w: torch.Tensor,
        block_size: List[int],
    ):
        if w.device.type == "xpu":
            from_hp_xpu(cls, w, block_size)
        elif w.device.type == "npu":
            from_hp_npu(cls, w, block_size)
     
           
implements = Int4PlainInt32Tensor.implements
implements_torch_function = Int4PlainInt32Tensor.implements_torch_function

@implements(aten.linear.default)
@implements_torch_function(torch.nn.functional.linear)
def _(func, types, args, kwargs):
    input_tensor, weight_tensor, bias = (
        args[0],
        args[1],
        args[2] if len(args) > 2 else None,
    )
    
	if input_tensor.device.type == "xpu":
    	return linear_xpu(input_tensor, weight_tensor, bias)
    elif input_tensor.device.type == "npu":
        return linear_npu(input_tensor, weight_tensor, bias)

Did I get that right? Happy to hear any thoughts or suggestions you might have!

jerryzh168 · 2025-10-21T17:55:33Z

You mean that we should keep the entry logic in quant_api.py unchanged:
and then handle different backend implementations in the from_hp and linear methods of Int4PlainInt32Tensor using simple if/else branches.

Yes that's correct

orangeH25 · 2025-10-22T06:34:41Z

You mean that we should keep the entry logic in quant_api.py unchanged:
and then handle different backend implementations in the from_hp and linear methods of Int4PlainInt32Tensor using simple if/else branches.

Yes that's correct

Got it, I will follow this approach, thanks!

orangeH25 · 2025-10-22T09:59:05Z

You mean that we should keep the entry logic in quant_api.py unchanged:
and then handle different backend implementations in the from_hp and linear methods of Int4PlainInt32Tensor using simple if/else branches.

Yes that's correct

Hi @jerryzh168 @jcaip , I’ve made those changes, please take a look, really appreciate it!

orangeH25 added 3 commits October 13, 2025 11:07

Add NPU (Ascend) backend support for INT4 weight-only quantization wo…

f3aefca

…rkflow

use torch.ops.npu prefix and drop redundant torch_npu import

68eea61

Merge branch 'pytorch:main' into quant/int4/wo/0

164435e

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 14, 2025

orangeH25 marked this pull request as draft October 14, 2025 11:44

fffrog reviewed Oct 14, 2025

View reviewed changes

Modify test file and update comments

06c77d1

orangeH25 marked this pull request as ready for review October 15, 2025 06:28

jcaip reviewed Oct 20, 2025

View reviewed changes

jerryzh168 reviewed Oct 20, 2025

View reviewed changes

Merge branch 'pytorch:main' into quant/int4/wo/0

498f052

add: merge NPU(Ascend) backend logic in Int4PlainInt32Tensor subclass

ea2aa7a

orangeH25 force-pushed the quant/int4/wo/0 branch from 7808297 to ea2aa7a Compare October 22, 2025 09:57

		@unittest.skipIf(torch_npu is None, "torch_npu is not available")
		@unittest.skipIf(not torch_npu.npu.is_available(), "NPU not available")

	@unittest.skipIf(torch_npu is None, "torch_npu is not available")
	@unittest.skipIf(not torch_npu.npu.is_available(), "NPU not available")
	@unittest.skipIf(torch.accelerator.current_accelerator(True).type == "npu" and torch.accelerator.is_available(), "NPU not available")

	f"torch_npu.npu_convert_weight_to_int4pack expects `int32` dtype"
	f"torch.ops.npu.npu_convert_weight_to_int4pack expects `int32` dtype"

Uh oh!

Add NPU (Ascend) backend support for INT4 weight-only quantization workflow #3172

Are you sure you want to change the base?

Add NPU (Ascend) backend support for INT4 weight-only quantization workflow #3172

Uh oh!

Conversation

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

Uh oh!

Summary

Environment

Files changed

Implementation Overview

Test Case

Uh oh!

pytorch-bot bot commented Oct 14, 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/3172

Uh oh!

fffrog Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

fffrog Oct 14, 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!

fffrog Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

fffrog Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

fffrog Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

orangeH25 commented Oct 15, 2025

Uh oh!

jcaip Oct 20, 2025

Choose a reason for hiding this comment

Uh oh!

jcaip left a comment

Choose a reason for hiding this comment

Uh oh!

jerryzh168 commented Oct 20, 2025

Uh oh!

jerryzh168 Oct 20, 2025

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

jerryzh168 commented Oct 21, 2025

Uh oh!

orangeH25 commented Oct 22, 2025

Uh oh!

orangeH25 commented Oct 22, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

orangeH25 commented Oct 14, 2025 •

edited

Loading

pytorch-bot bot commented Oct 14, 2025 •

edited

Loading

fffrog Oct 14, 2025 •

edited

Loading

orangeH25 commented Oct 21, 2025 •

edited

Loading