Repository: pytorch
Issue: RFC: Integration of KleidiAI 4-Bit MatMul Kernels into PyTorch
Labels: oncall: quantization,triaged,matrix multiplication

# RFC: Integration of KleidiAI 4-Bit MatMul Kernels into PyTorch

**Hardware Platform**: aarch64

## Motivation
1. Add support for [KleidiAI](https://gitlab.arm.com/kleidi/kleidiai) 4-bit matrix multiplication kernels to PyTorch.
2. Expose these 4-bit matmul kernels as PyTorch operators to accelerate ML models.

**Pytorch Integration PR**:  [#134124](https://github.com/pytorch/pytorch/pull/134124)
**torch/ao Quantization Change**:  https://github.com/ng-05/ao/commit/cbcf915535a0f3f244a21853caf6ef0b321d7ede

## Overview of 4-Bit MatMul Operations
The 4-bit matmul process consists of two main operations:
1. **Packing**: Quantized weights (4-bit), floating-point scales and biases are combined into a single `packed_weight` buffer.
2. **Execution**: Perform matrix multiplication on FP32 input using the `packed_weight` buffer.

## Current Target Operations
We aim to enhance two existing PyTorch operations:
1. **`_convert_weight_to_int4pack_cpu`**: This operation is responsible for converting weights to a packed format.
2. **`_weight_int4pack_mm_cpu`**: This operation performs the matrix multiplication using the packed weights.

### Issues Identified

1. **Modification of Operation Signatures**:
   - **For `_convert_weight_to_int4pack_cpu`**:
     - **Reason**: The KleidiAI packing kernel requires scales, bias, and 4-bit quantized weights to create a single linear buffer.
     - [Change Link](https://github.com/pytorch/pytorch/pull/134124/files#diff-db2b91c51cffd1fba933933ac20de256f5f0ced095890a1a49a17d73037001a2R3434)
   - **For `_weight_int4pack_mm_cpu`**:
     - **Reason**: The packing operation loses the shape information of the original weights (N, K). The `packed_weight` shape is 1 dimensional and accounts for scales, bias , weights. We need the row count (N) to perform matmul correctly.
     - [Change Link](https://github.com/pytorch/pytorch/pull/134124/files#diff-db2b91c51cffd1fba933933ac20de256f5f0ced095890a1a49a17d73037001a2R3512)

2. **Support for Channelwise and Groupwise Quantization**:
   - **Need**: We require a way to pass group size information to the underlying kernel for channelwise and groupwise quantized matmuls. Currently, a group size of 0 is used to indicate channelwise quantization and 32 to indicate 32 groupwise quantization

3. **Data Type Handling in Operations**:
   - **Issue**: The existing `_weight_int4pack_mm_cpu` operator performs multiplication and accumulation in FP32/BF16. In contrast, our kernels dynamically quanitzes fp32 input to INT8, use INT8 for multiplication and accumulate results in FP32. This might introduce noticeable (but within acceptable error range [_mean error : **0.0064**_]) accuracy changes for the same operation across platforms.
   - **Proposed Solution**: Introduce a context manager to handle low-bit operations explicitly within the same PyTorch operator, conveying the datatype difference to the user. Someting similar and as simple as [torch._C._set_cpu_allow_fp16_reduced_precision_reduction](https://github.com/pytorch/pytorch/blob/f69bf005f7b3be70b9a38ba663e6bc825dfe7ffd/torch/csrc/Module.cpp#L1487) 

## Initial Approach:
   - We initially introduced new PyTorch operators due to above mentioned issues and the lack of support for 4-bit matmul operations using 8-bit intrinsics. However, this approach was quickly flagged by the Meta pytorch team.

## Why Kleidi in pytorch?
  - Kleidi will expand beyond lowbit matmul kernels, we would like to add support for Kleidi kernels directly in existing pytorch operators so that optimizations can be leveraged directly for ML Models

## Proposed Path Forward
To resolve the above issues and successfully integrate our kernels into `_convert_weight_to_int4pack_cpu` and `_weight_int4pack_mm_cpu`, we seek suggestions and collaboration on the following:

- Adjusting the operation signatures to accommodate the necessary parameters.
- Finding a method to pass group size information for quantization.
- Ensuring proper handling of data types during the matmul process.
- Considering the implications of the context manager approach.



## E2E Flow
Our torch ao quantizer implementation can directly replace GPTFast block in below diagram
[
![Kleidiai_Pytorch_flow_2](https://github.com/user-attachments/assets/19ce9592-c17a-4c9d-81e4-90197407e48f)
](url)

Existing 4 bit matmul kernel fucntionality in `_weight_int4pack_mm_cpu` operartion : 
```
FP32 Input                INT4 Weights     FP32 Scales
[32.1, 15.7, -3.9]        [3, -2, 1]     [0.3, 0.2, 0.15]
       |                      |            / 
       |                      |           / 
       |                Dequantization   / 
       |                      |         /
       |                      v        /
       |             FP32 Dequantized Weights
       |            [0.375, -0.25, 0.125]
       |                      |
       |                      |
       +----------------------+
                  |
                  |
             Matrix Multiplication
                  |
                  v
         FP32 Output
    [12.09375, -3.9375, -0.48125]
```
KleidiAI 4 bit matmul kernel fucntionality in `_weight_int4pack_mm_cpu` operartion : 
```
FP32 Input                Packed Buffer ( INT4 Weights + scales + Bias )
[32.1, 15.7, -3.9]        [3, -2, 1, 0]    [8.25, 3.5]  [1.2, 5.5]
       |                       |
       |                       |
    Quantize              Dequantize
   to INT8                 to INT8
       |                       |
       v                       v
INT8 Input              INT8 Weights
 [64, 31, -8]           [24, -16, 8, 0]
       \                     /
        \                   /
         \                 /
          \               /
           \             /
            \           /
             \         /
           INT8 Matrix Multiplication
                   |

                   |
                   v
 FP32 Dequantized and Accumulate in FP32
         [12.0, -3.875, 1.5, 0]

                   |
                   v
             FP32 Final Output
            [9.625, -3.875, 1.5]
```

Your feedback and suggestions are highly welcome!

cc: @malfet @digantdesai @jgong5 @sanchitintel @cfRod @milpuz01 

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim