search

flashinfer-ai / flashinfer

FlashInfer: Kernel Library for LLM Serving
https://flashinfer.ai
Apache License 2.0
1.14k stars 103 forks source link

How to use low-bit KV Cache in flashinfer? #125

Open zhaoyang-star opened 7 months ago

zhaoyang-star commented 7 months ago

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 commented 7 months ago

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

zhaoyang-star commented 7 months ago

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 commented 7 months ago

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

yzh119 commented 7 months ago

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 commented 6 months ago

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 commented 5 months ago

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.

SherrySwift commented 5 days ago

Hi, is there any plan to integrate the 4-bit fused dequantize+attention operators proposed in Atom into FlashInfer? Looking forward for this new feature.