[RFC] A Proposal to Support More Accelerators for TorchAO

[orangeH25]

Summary

We propose a modification plan about extending TorchAO to support more accelerators easier.

Key points:

• Generic quantization path – a universal implementation available for all backends.
  • Quantization stage: Lower performance priority. Implemented using standard torch.ops composition to ensure generality and compatibility.
  • Inference stage: Performance-sensitive. Higher priority. Uses PyTorch’s operator registration/dispatch mechanism to balance performance and portability.
• Specific hardware path – for backends that require custom, non-generic implementations.
  • Example: The current XPU INT4 weight quantization branch, which uses a backend-specific implementation.

Current Status

Take INT4 weight quantization as an example.

Current state:

• FORMAT: PLAIN → Intended as a multi-device generic format, but currently restricted to CUDA due to FBGEMM limitations.
• FORMAT: PLAIN_INT32 → A custom quantization format path specific to the XPU backend.

Proposed Implementation

We continue with INT4 weight quantization as the running example.

Generic Quantization Path

Quantization Stage

The quantization stage mainly consists of two types of operations: quantization and packing.

These operations are implemented in AO with standard Torch operators to guarantee portability:

def int4_row_quantize_zp(
    x: torch.Tensor,
    group_size: int = 128,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # quantize operator, implemented with standard torch ops
    ...

def pack_int4(x: torch.Tensor) -> torch.Tensor:
    # pack operator, implemented with standard torch ops
    ...	

Inference Stage

Take bf16i4bf16_rowwise Operator as example:

AO-side Implementation (TorchAO Repo)

1. Register the backend-specific operator signature through PyTorch’s operator registration system:

TORCH_LIBRARY(ao, m) {
    m.def("bf16i4bf16_rowwise(Tensor input, Tensor weight, Tensor scale, Tensor zero_point) -> Tensor");
}

2. Provide a generic fallback implementation using standard Torch operators:

at::Tensor bf16i4bf16_rowwise(
    at::Tensor X, // BF16
    at::Tensor W, // INT4
    at::Tensor w_scale_group,
    at::Tensor w_zero_group
) {
	// Performs BF16 × INT4 computation  
	// Uses standard PyTorch operators to enable a generic fallback implementation
    ...
}

Device-side Implementation (Device Repo)

• Register a device-specific kernel with AO. During inference, this implementation will be prioritized:

TORCH_LIBRARY_IMPL(ao, device, m) {
    m.impl("bf16i4bf16_rowwise", bf16i4bf16_rowwise);
}

---

When this operator is invoked in AO, the dispatcher selects the device-specific implementation if available. Otherwise, it falls back to the generic implementation:

res = torch.ops.ao.bf16i4bf16_rowwise(
    input_tensor,
    weight_tensor.qdata,
    weight_tensor.scale,
    weight_tensor.zero_point,
)

Specific hardware Path

For devices requiring custom, non-generic implementations, a separate branch can be created. The existing XPU implementation serves as a reference:

elif int4_packing_format == Int4PackingFormat.PLAIN_INT32:
	# device-specific path for PLAIN_INT32 format
	new_weight = Int4PlainInt32Tensor.from_hp(
    	weight,
        block_size,
    )
    return new_weight

Benefits

Improved extensibility: By using standard PyTorch operators, hardware backends can leverage quantization support while ensuring wide portability.

Note on Scope

This RFC proposes a general extensible quantization framework. INT4 weight quantization is used as an example, but the approach is generic. The initial implementation may focus on INT4, with later extensions to INT8, FLOAT8, or other quantization schemes.

[orangeH25]

Hi @jerryzh168 @huydhn , Could you take a look when you have a chance? Thanks!

[jcaip]

Hey @orangeH25 can you share more details on your proposal?

What accelerators do you plan to support?
Are there gaps in the current tensor subclass flow? Why can't we use our existing dispatch mechanism to specify accelerator specific behavior?

[orangeH25]

Hey @jcaip

What accelerators do you plan to support?

Our motivation is to enable quantization support for the Ascend NPU (integrated into PyTorch as an out-of-tree backend via PrivateUse1). To facilitate integration of other PyTorch backends with TorchAO while minimizing code changes in TorchAO itself, we propose this RFC using int4 as an example.

Are there gaps in the current tensor subclass flow?

Yes, integrating a new device through a Tensor subclass is also a good approach. However, this approach has key limitations: some common logic cannot be shared across PyTorch backends, and more importantly, each PyTorch backends would require changes to TorchAO’s codebase, which would add maintenance overhead in the long run.

Why can't we use our existing dispatch mechanism to specify accelerator specific behavior?

Sorry, could you elaborate a bit more on this? If you are referring to the dispatch capability of the operator torch.ops.fbgemm.bf16i4bf16_rowwise, the main issue is that many of FBGEMM’s features are currently only supported on GPU.

cc @FFFrog

[jcaip]

Sorry, I still don't think I quite "get" this. Let me bring it up at the team meeting and see what the others say. I guess my question here is, why not do something similar to our existing XPU Int4 implementation here

Yes, integrating a new device through a Tensor subclass is also a good approach. However, this approach has key limitations: some common logic cannot be shared across PyTorch backends

You could share logic here via class inheritance no?

, and more importantly, each PyTorch backends would require changes to TorchAO’s codebase, which would add maintenance overhead in the long run.

Are you proposing a way to be able to add new backend without needing to modify the TorchAO codebase?

[jerryzh168]

could you expand more on Specific hardware Path, what is this referring to

[orangeH25]

Hey @jcaip @jerryzh168

Take an example:

Background: A new backend aims to support AO’s int4 PLAINI format weight quantization.

Approach 1

Similar to the existing implementation

1. Add a new branch in _int4_weight_only_quantize_tensor:

# torchao/quantization/quant_api.py 

elif int4_packing_format == Int4PackingFormat.NEW:
    new_weight = Int4NewTensor.from_hp(
        weight,
        block_size,
    )
    return new_weight
else:
    raise ValueError(f"Unsupported int4 packing format: {int4_packing_format}")

2. Create the corresponding Int4NewTensor subclass and implement backend-specific logic:

# torchao/quantization/quantize_/workflows/int4/int4_new_tensor.py

class Int4NewTensor(TorchAOBaseTensor):
    ...
    @classmethod
    def from_hp(
        cls,
        w: torch.Tensor,
        block_size: List[int],
    ):
        ...

implements = Int4NewTensor.implements

@implements([torch.nn.functional.linear, aten.linear.default])
def _(func, types, args, kwargs):
    ...
    
...

Approach 2

Based on our proposed design

we either extend the functionality of an existing tensor subclass or introduce a new generalized tensor subclass.

For example, we extend the existing PLAIN → Int4Tensor path to make it more general.

1. General entry point（accessible to all PyTorch backends）

# torchao/quantization/quant_api.py 

elif int4_packing_format == Int4PackingFormat.PLAIN:
    new_weight = Int4Tensor.from_hp(
        weight,
        block_size,
    )
    return new_weight

2. Quantization and packing logic (already implemented with a generic implementation, as described in the RFC)

# torchao/quantization/quantize_/workflows/int4/int4_tensor.py

# Int4Tensor.from_hp
wq, scale, zero_point = int4_row_quantize_zp(w, group_size)	
wq = pack_int4(wq)

3. Inference stage, for operator overrides such as:

# torchao/quantization/quantize_/workflows/int4/int4_tensor.py
@implements([torch.nn.functional.linear, aten.linear.default])
def _(func, types, args, kwargs):
    ...
    res = torch.ops.ao.bf16i4bf16_rowwise(
        input_tensor,
        weight_tensor.qdata,
        weight_tensor.scale,
        weight_tensor.zero_point,
    )
    ...

As described in the RFC, we create ao namespace and use dispatch mechanism, so that different PyTorch backends can register their own backend-specific implementations in their respective repositories.

At the same time, AO provides a generic fallback implementation to ensure functionality when no backend-specific kernel is available.

---

We recommend this second approach because it provides a unified path that works across all PyTorch backends. This is similar to AO’s existing INT8 weight quantization, where most PyTorch backends can run without modification.

[jcaip]

This reminds me a lot of the _layout parameter we are currently moving away from, see here. For example for Int4WeightOnly we support a couple different Layouts - Tinygemm, Marlin Sparse, Marlin QQQ.

I think the opinion of the team was that this is a little confusing to develop and maintain, and moving each layout to it's own tensor is a better approach. cc @jerryzh168 if I'm missing some of the reasoning here.

We recommend this second approach because it provides a unified path that works across all PyTorch backends. This is similar to AO’s existing INT8 weight quantization, where most PyTorch backends can run without modification.

Are there functionality gaps that differ between the two approaches? It doesn't seem like there are any to me. Given that this would be a pretty significant refactor, and that it's something we're currently moving away from, I don't think we would be interested in this proposal.

We'd still be interested in accepting backend support for AO different accelerators though :)

[orangeH25]

Are there functionality gaps that differ between the two approaches?

Yes, there are differences. In Int4WeightOnly v1, the _layout parameter is used to trigger additional logic, with backend-specific implementations placed inside the corresponding class, i.e., within the AO repo.

In contrast, our proposal builds on the structure of Int4WeightOnly v2, where a common tensor subclass is provided. Backend-specific optimized kernels would then be implemented and registered within each backend’s own repo. While this approach may offer slightly less flexibility, it provides a unified and consistent integration path for all backends.

[orangeH25]

We'd still be interested in accepting backend support for AO different accelerators though

Hey @jcaip @jerryzh168 , we highly value the current Int4WeightOnly v2 scheme, which leaves more flexibility for different backends and makes it easier for them to integrate according to their own characteristics.

At present, the mainstream way to add a new backend in PyTorch is to integrate it as an out-of-tree backend via PrivateUse1, with Ascend NPU being a typical example. As the ecosystem continues to grow, we believe more and more backends will adopt this approach, making such scenarios increasingly common.

We would like to explore whether Ascend NPU could be integrated into AO by following the approach used for XPU, so that it can support AO’s quantization features.

[jcaip]

In contrast, our proposal builds on the structure of Int4WeightOnly v2, where a common tensor subclass is provided. Backend-specific optimized kernels would then be implemented and registered within each backend’s own repo. While this approach may offer slightly less flexibility, it provides a unified and consistent integration path for all backends.

I don't think this is the case, if we rely on backend for dispatch, then we'll still need to have multiple subclasses to handle the case when we have multiple packing formats (Gemlite, Marlin, TensorCoreTiledLayout) for a shared backend (cuda).

Additionally, these layouts all have custom packing layouts, which can't be shared. Do you know how the Ascend NPU expects the weights to be backed for Int4WeightOnly quantization? Is that consistent with one, if any of our existing int4 packing formats?

We would like to explore whether Ascend NPU could be integrated into AO by following the approach used for XPU, so that it can support AO’s quantization features.

Yes, please we'd be happy to adding add NPU support the same way we enabled XPU support. But I don't think the general solution proposed here is the right path for TorchAO.

[jerryzh168]

It seems that the proposal is to extend the device coverage for the final quantized operator e.g. bf16i4bf16_rowwise directly, while sharing all the other integration path. I think it makes sense if:

• the packing format that's picked, e.g. plain packing for int4, is also optimal for the target device, e.g. Ascend NPU (we have: device = "ascend_npu")
• the quantization is also done in a way that's compatible with the desired packing format, this includes the slight differences of choose_qparams and quantize op, and the dtype for packed weights etc. e.g. XPU want these settings:

  ao/torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor.py

  Lines 102 to 108 in 7690612

  	mapping_type = MappingType.ASYMMETRIC
  	target_dtype = torch.int32
  	quant_min = 0
  	quant_max = 15
  	eps = 1e-6
  	scale_dtype = None
  	zero_point_dtype = torch.int32

If this is the case, it's already a supported use case and fit into our current design in v2 I think.

[orangeH25]

Hey @jcaip @jerryzh168 , we really appreciate your reply and suggestions.

Yes, please we'd be happy to adding add NPU support the same way we enabled XPU support. But I don't think the general solution proposed here is the right path for TorchAO.

Thank you for the suggestion. We will proceed with the integration following the XPU approach, and once the PR is ready, we’d appreciate your review.

Do you know how the Ascend NPU expects the weights to be backed for Int4WeightOnly quantization? Is that consistent with one, if any of our existing int4 packing formats?

If we adopt the same implementation path as XPU, our current thinking is that the packing format expected by Ascend NPU is consistent to XPU’s PLAIN_INT32

---

By the way, if we implement a tensor subclass with an int32 packing format similar to XPU (temporarily named Int4PlainInt32TensorNPU), we could place it in the entry as follows:

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

elif int4_packing_format == Int4PackingFormat.PLAIN_INT32_NPU: 
    new_weight = Int4PlainInt32TensorNPU.from_hp(
        weight,
        block_size,
    )
    return new_weight

or

elif int4_packing_format == Int4PackingFormat.PLAIN_INT32:
    if device == "xpu":
        new_weight = Int4PlainInt32Tensor.from_hp(
            weight,
            block_size,
        )
    elif device == "npu":
        new_weight = Int4PlainInt32TensorNPU.from_hp(
            weight,
            block_size,
        )
    return new_weight

What’s your opinion on this?