Q Scaling factor for FP8 KV Cache ?

[Syst3m1cAn0maly]

When using a FP8 quantized model with FP8 KV cache scales I get this warning in newer versions of vLLM :
WARNING 03-25 22:34:30 [kv_cache.py:82] Checkpoint does not provide a q scaling factor. Setting it to k_scale. This only matters for the flash-attn backend.

I indeed use Flash Attention 3.

How can I provide the Q scaling factor to vLLM ?

I use this recipe :

recipe = """

quant_stage:
quant_modifiers:
QuantizationModifier:
ignore: ["re:.*lm_head"]
config_groups:
group_0:
weights:
num_bits: 8
type: float
strategy: tensor
dynamic: false
symmetric: true
input_activations:
num_bits: 8
type: float
strategy: tensor
dynamic: false
symmetric: true
targets: ["Linear"]
kv_cache_scheme:
num_bits: 8
type: float
strategy: tensor
dynamic: false
symmetric: true
"""

[brian-dellabetta]

Hi @Syst3m1cAn0maly , this appears to be more of an issue for vllm than llm-compressor. Is this comment relevant to you?

[Syst3m1cAn0maly]

@brian-dellabetta It does not seem to be related to the comment you mentioned as Flash Attention 3 supports FP8 KV Cache contrary to the older version mentioned.

It seems to me that Flash Attention 3 expects an additional scale for Q which is not computed in llm-compressor in the "kv_cache_scheme" part of the recipe. That's what I get from the warning but I might be mistaken.

[brian-dellabetta]

Hi @Syst3m1cAn0maly , yeah it looks like you are right. The warning is hit here because there is no q_scale Parameter, only k_scale and v_scale. Indeed, it looks like those are the only parameters set during llm-compressor's KVCache Quantization, as noted here.

The ramifications of this are unclear to me, does the model perform poorly?

[Syst3m1cAn0maly]

It does not seem to perform poorly but it might lose some precision compared to FlashInfer.
It is quite hard to notice, so I'm not sure how it could affect the responses.

I merely want to make sure I get the best precision for the FP8 KV cache and since I saw the warning I wanted clarifications to better understand the impact.

[brian-dellabetta]

Yeah, makes sense. I will ask one of the vllm maintainers to take a look, if we should add logic to llm-compressor to attach q_scales as part of the kv cache quantization.

[brian-dellabetta]

Hey @Syst3m1cAn0maly , I got some feedback from one of the vllm maintainers:

how we implemented fp8 kv cache in flash attention 3 is by leveraging the fp8 attention implementation, thus the need to have scales for query as well. So I would still recommend treating kv cache quantization and attention quantization separately. This is just an artifact of the kernel support we have in FA3

Our research team is evaluating the accuracy of this, though I am unsure of its priority. From llm-compressor side, it wouldn't be as easy as hoped to add q_scale. We subclass transformers.DynamicCache to implement kv cache quantization, so we cannot change the function signature of update to include query_states.

I am not super familiar with the specifics of flash attention and why this is only an issue for FA3. If you experience degradation or have further questions, consider creating an issue on vllm. Sorry I can't be of more help.

[Syst3m1cAn0maly]

Thank you very much @brian-dellabetta for the investigation

[aman-17]

Hey @brian-dellabetta, any update on this??

[brian-dellabetta]

Hi @aman-17, hi @Syst3m1cAn0maly -- I asked around, and @eldarkurtic is currently investigating this. My understanding is that this is an artifact of the system defaulting (at least in some circumstances) to Flash Attention 3, and is not expected to cause downstream performance issues. But it is more of a vllm concern than llm-compressor, so I am not an expert.

cc @eldarkurtic , is there any additional information we can provide users on this?

[eldarkurtic]

Hey folks, I am actively working on this right now so I can provide more details. The reason you are seeing this warning is because the current version of LLM-Compressor doesn't support calculating scales for queries (q_scale), only for k_scale and v_scale. FlashAttention-3 expects to get scales for all three (query, keys, values), so due to the lack of q_scale in LLM-Compressor produced checkpoint, vllm will use a scale of 1.0 behind the scenes (at least with FlashAttention backend, I haven't checked other backends).

There are two active on-going items which will be relevant for getting the best accuracy with quantized KV-cache:

1. Enabling calculation of q_scale in LLM-Compressor so we can use something better than 1.0
2. The existing versions of LLM-Compressor and vLLM support only per-tensor scales for Q, K, V which is not the best choice for preserving accuracy. We have a first prototype which enables per-head scales for Q, K, V and we will try to release it over the next few weeks.