Feature Request: Implement Static Cache and Quantization Techniques in CTranslate2

[Jiltseb]

@minhthuc2502 @alexlnkp

Description
What type of cache is currently implemented in CTranslate2? Is it static or dynamic? Could we achieve a speed-up if the cache implementation is changed for the decoder in encoder-decoder models?

Also, it would be great to implement recent popular quantization techniques such as [HQQ] (https://github.com/mobiusml/hqq) in ctranslate2 format.

Motivation
Given that a static cache (see this PR) can significantly speed up processing in PyTorch encoder-decoder models via torch compilation, can we enable this in CTranslate2? This enhancement can improve decoding speed for projects utilizing CTranslate2 models, such as Faster Whisper.

References
Speed-up achieved for PyTorch-based Whisper: Blog Post

Benefits
Implementing static caching and recent quantization techniques in CTranslate2 could lead to significant performance improvements in model decoding speeds and efficiency.

Thank you for considering this feature request!

[minhthuc2502]

Hello, thank you for your information. I will take a look to know how hqq works. Otherwise, the cache's implementation in Ctranslate2 is reallocated depending on the length of sequence. I'm not sure if the static cache (as my understanding, the cache will be pre-allocated before with max size) can speed up much with Ctranslate2, it requires some benchmark to confirm it. Some changes in design is needed.

[Jiltseb]

Hello, thank you for your information. I will take a look to know how hqq works. Otherwise, the cache's implementation in Ctranslate2 is reallocated depending on the length of sequence. I'm not sure if the static cache (as my understanding, the cache will be pre-allocated before with max size) can speed up much with Ctranslate2, it requires some benchmark to confirm it. Some changes in design is needed.

@minhthuc2502 Thanks. This blog post should help with HQQ: https://mobiusml.github.io/hqq_blog/

[minhthuc2502]

Hello, I tried to implement the HQQ for 4 bit quantization, it works but I think it has to combine low-bit matmul kernel to speed up the inference, Do you have any reference for this?

[mobicham]

@minhthuc2502 we use the int4mm kernel from torchao: https://github.com/mobiusml/hqq/blob/master/hqq/backends/torchao.py

[minhthuc2502]

As I understand it, to use int4mm, you have to convert HQQ quant to the format accepted by int4mm (w_q and scale,,,). Do you think it will reduce the performance? Similar to using int4mm kernel, It looks like when we dequantize the weight and make the gemm, and the performance is not as expected. I see you have very good performance with torch.compile but in Ctranslate2, we don't have this kind of solution. I will investigate more about others solutions and the kernel along with them.

[mobicham]

The conversion is only done once, different int4 kernels require different input formats, so that's why we do it via patching so we can support many backends, not just the torchao int4 kernel.

The int4mm kernel should be faster than fp16 matmul with or without torch.compile. However, it is a gemv kernel optimized for decoding only, so the prefill phase with this kernel is actually slower. That's why we don't use it in the encoder, and only use it in the decoding phase. Unfortunately, there's no way so far to efficiently dequantize() and do fp16 matmul in the prefill phase which should be faster. But decoding one-token at a time with the int4mm is def faster.

There are other options by the way, like https://github.com/microsoft/BitBLAS/ , they also support A16W2 matmul which should be even faster for larger models.

[Jiltseb]

@minhthuc2502 Did you figure out how to speed up the HQQ implementation in ctranslate2? This will be a useful add-on for large E-D models.

[minhthuc2502]

I tried to implement HQQ, only quantization, but I have not do the benchmark yet. The thing which prevented me to make it work correctly is patching the HQQ format. As I understand I will patch the HQQ format in the conversion time and then do the inference with the new format. BTW, I'll try to implement it.

[minhthuc2502]

I tried hqq quant in this PR . When I use int4mm, the quality is not good and it becomes worse with long prompt. I tested only with decoder. Do you have any idea if I did any mistake in the conversion.

I tried using hqq quant first and then converting directly to torch quant format to have scales and zeros and weights. Then before the inference, I will convert the weight to int4pack and then use the int4mm for matmul.

[mobicham]

@minhthuc2502 you have to quantize with hqq. The link you shared is just doing RTN quantization, which will give bad quality especially at lower bits.

[minhthuc2502]

I followed this to add the hqq https://github.com/mobiusml/hqq/blob/master/hqq/core/quantize.py. Did I miss something?

[mobicham]

@minhthuc2502 yes, it's not performing optimization: https://github.com/mobiusml/hqq/blob/master/hqq/core/quantize.py#L115-L122 , which is the actual hqq algo: https://github.com/mobiusml/hqq/blob/master/hqq/core/optimize.py#L194-L243
Basically, you get an initial estimate of the quantized weights/scale/zero (which is what you did), but then you fine-tune them via alternate minimization.
Any reason why you re-implement things in numpy ? You could simply call HQQLinear from the hqq lib, it will do everything for you. Doing this on cpu will be very slow, better do it on the GPU, example below:

from hqq.core.quantize import HQQLinear, BaseQuantizeConfig
from hqq.backends.torchao import patch_hqq_to_aoint4
quant_config = BaseQuantizeConfig(nbits=4, group_size=64, quant_zero=False, quant_scale=False, axis=1)
hqq_ao_layer = patch_hqq_to_aoint4(HQQLinear(your_linear_layer, quant_config=quant_config, compute_dtype=torch.bfloat16, device='cuda'), None)

out = hqq_ao_layer.forward(x)

[minhthuc2502]

Thank you for your suggestion. I just use directly the HQQLinear and the quality is better due to the optimization. See in the PR. When I tested with long prompt (tried to summarize a long text), it did not work good, there is some repetitions or worse. What do you think about it? Is it normal due to the lower bit quantization?