How to use low-bit KV Cache in flashinfer?

[zhaoyang-star]

From the blog I noticed that FlashInfer implements low-precision attention kernels so that we can achieve nearly linear speedup to the compression ratio (~4x for 4bit, ~2x for 8bit). This feature is great! and I try to use it. But there is no demo or toy code about how to use it. Could you please share more details about it?

[yzh119]

I haven't exposed low-bit KV-Cache in PyTorch APIs (they are available in C++ APIs), will do it tmr :)

[zhaoyang-star]

Glad to hear that! Cannot wait to try it out. I think quantizing KV Cache from float16/bfloat16 to 4-bits will need calibration. It will be better if the feature released with demo and benchmark results (latency, throughput or accuracy).

BTW, there is already someone trying to port flashinfer to vLLM (see #2772) to boost decode phase. I also ported FlashAttention to vLLM (see #2744) and plan to benchmark FA and flashinfer in vLLM framwork.

[yzh119]

Thanks for letting me know, it's interesting to see that FlashAttention starts supporting paged kv-cache.

[yzh119]

It will be better if the feature released with demo and benchmark results (latency, throughput or accuracy).

You can check our manuscript: Atom: Low-bit Quantization for Efficient and Accurate LLM Serving.

[yzh119]

PyTorch APIs for fp8 kv-cache are exposed in #156 .

I'm finalizing the int4/int8 fused-dequant attention kernels with some optimizations such as fast int4/int8-to-float16 conversions. I expect to merge these changes by this Thursday.

[zhyncs]

PyTorch APIs for fp8 kv-cache are exposed in #156 .

I'm finalizing the int4/int8 fused-dequant attention kernels with some optimizations such as fast int4/int8-to-float16 conversions. I expect to merge these changes by this Thursday.

Hi @yzh119 As mentioned in https://flashinfer.ai/2024/02/02/introduce-flashinfer.html.

Our next release will include the 4-bit fused dequantize+attention operators proposed in Atom and LoRA operators used in Punica.

When is Atom quantization expected to be fully integrated into FlashInfer? Is there a detailed timeline available? Thanks.