Question about KV Cache quantization

[SherrySwift]

Hi, thanks for your great work!
I have a small question about KV Cache quantization. Did you use pagedattention to accelerate KV Cache 4-bit quantization? If so, where is the corresponding cuda kernel code? Thank you.

[happierpig]

Hi @SherrySwift ,

Yes, we use page attention with low-bit quantization to speedup memory loading time. INT4 tokens are loaded and decoded into FP16, which still uses FP16 arithmetic calculations. Check quantization code here: https://github.com/efeslab/Atom/blob/main/kernels/include/flashinfer/quantization.cuh#L10

[SherrySwift]

Hi @happierpig ,
Thanks for your quick reply, I have another question.
According to

Atom/kernels/include/flashinfer/quantization.cuh

Line 70 in 7e3618b

constexpr size_t PACK_NUM = 8;

, 8×4bit data are packed into one int32 data.
But in

Atom/e2e/punica-atom/punica/utils/kvcache.py

Line 19 in 7e3618b

dtype=torch.uint8,

, it seems that 4bit KV Cache are stored in int8 type rather than int32 type.
Feel a little bit confused about this.

[happierpig]

Sure! We use int8 data type to pack 2 int4 elements since the minimum addressable data size on GPU is byte.