Repository: pytorch
Issue: Faster Pytorch dequantize() + matmul for quantized models
Labels: oncall: quantization

### Fast Pytorch dequantize() + matmul

I would like to open the discussion about faster inference with quantized models using pure Pytorch calls. 

As you know, quantization is extremely important to run large models like LLMs with limited GPU memory. This is especially important for the open-source community. Making these models more accessible is a huge step forward. 

The quality of quantized models is improving at a fast pace. With the recent releases of [HQQ](https://github.com/mobiusml/hqq/) and [Quip#](https://github.com/Cornell-RelaxML/quip-sharp/). We have now reached a point where 4-bit and even 2-bit models work fairly well and this is surely gonna improve in the upcoming months. 

Now the only issue: **inference speed**. While some libraries like llama.cpp are able to achieve _faster_ inference with quantized models, a naive implementation in Pytorch can be 2x slower compared to fp16.

### Alternatives
#### Custom Cuda/Triton kernels
Most of the methods like BNB/GPTQ/AWQ/Quip etc. implement custom CUDA/Triton kernels to do the dequantize() + matmul step. These kernels require a very low-level implementation and that can only be done efficiently in special cases. Moreover, these kernels are usually optimized for newer Nvidia chips like A100s and it's not possible to use this approach for other hardware like CPU, Intel's Arc, etc. In fact, some Triton kernels optimized for A100 may run _slower_ on older GPUs. 

#### Naive Pytorch Implementation
With a pure Pytorch implementation,  this is basically implemented as follows:
```python
def forward(self, x):
	W_est = self.dequantize()
	out   = torch.matmul(x, W_est.t())
	if(self.bias!=None): out += self.bias
	return out
```
The issue is that this is relatively slow due to the ```dequantize()``` call. 

#### Torch Compile
This actually does make things faster with dynamo compile but:
- it's still not that fast.
- this actually creates issues with async calls for LLM token streaming for example.

#### Aten C++ Backend
I re-implemented the whole dequantization logic in Aten: https://github.com/mobiusml/hqq/blob/master/hqq/kernels/hqq_aten.cpp
but the dynamo compile option is faster. 

I found this new function that could be used in Aten but there's no documentation on how to use it: https://pytorch.org/cppdocs/api/function_namespaceat_1adeda9630914278ac02d7fd758da19e3d.html

### Ideas
### Efficient Native Pytorch bit-unpacking
Really the main speed bottleneck here is the bit-unpacking. I think we really need a native well-implemented bit-packing/unpacking function in Pytorch. The current hack to make low-bit tensors is to store them in uint8/uint32. The bit-unpacking operation requires extra copies and the operations when implemented via torch.cat for example are not optimized because they are done in serial. For 3-bit it's actually much worse because we need to use uint32 instead of uint8 for bit-packing which requires 10 mask+shfit+copy + another slicing operation because the weight shapes are not a multiple of 3.

In an ideal world, it would be awesome to have native Pytorch **uint4**, **uint3**, **uint2** dtypes that work via casting:
```python
t_2bit_as_8bit = torch.tensor([0, 1, 2, 3], dtype=torch.uint8)
t_2bit_as_2bit = t_2bit_as_8bit.to(torch.uint2) # actual 2-bit tensor
```
I know that `quint4x2` already exists but, correct if I am wrong, it's not possible now to do uint8 <> quint4x2 casting. I was not able to use it directly.

I think by just having a more efficient implementation of bit-unpacking and use Dynamo for compilation we can potentially have something just as fast as fp16, if not faster. The implementation can be simply doing the same mask+shift+copy,  but in a more efficient way. For example, I noticed that replacing the torch.cat operation by first allocating an empty array then copying the chunks can be quite faster. 

### Fused Dequantize() + matmul() kernels 
As I have mentioned earlier, there's one here: https://pytorch.org/cppdocs/api/function_namespaceat_1adeda9630914278ac02d7fd758da19e3d.html
I am curious to know if you plan to have something similar for 2 and 3 bit. 
Also, how much is it faster compared to fp16?
This solution will not work for codebook-based approaches though.

### Conclusion
I think it would be really great to open the discussion about this topic. Happy to know your roadmap regarding accelerating quantized models with Pytorch. Thanks! 


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