Analysis of Why KV Cache Quantization Logic in kv_cache_quant.py Is Not Taking Effect

[MELANCHOLY828]

I have a question mainly about the implementation inside kv_cache_quant.py. I noticed that in quantized_forward, there is a condition checking if isinstance(outputs, tuple) and len(outputs) >= 3:. Only if this holds, the KV quantization path is taken. However, when I actually run exp_niah.py, I found that this condition is not met — meaning the KV quantization doesn’t take effect.

[jagmarques]

The quantization logic not taking effect is a common issue with KV cache quantization implementations — the quant/dequant hooks often get bypassed if the attention backend switches to a fused kernel path that doesn't call the Python-level hooks.

A few things to check in kv_cache_quant.py:

1. Hook registration: verify the hooks are attached to the right module. In transformer implementations, KV cache is often written inside _attn or self_attn.forward — if quantization hooks are on the outer module, they may not intercept the actual tensor writes.

2. Fused attention bypass: if the model uses torch.nn.functional.scaled_dot_product_attention with is_causal=True, it may take a Flash/cuDNN path that bypasses custom hooks entirely. Force attn_implementation='eager' to confirm.

3. In-place vs copy: some KV cache implementations update the cache in-place — if your quantization returns a new tensor rather than modifying in-place, the original cache tensor won't be updated.

For reference, NexusQuant handles this by intercepting at the attention module level with a context manager (with nexusquant(model):) rather than registering hooks, which sidesteps the fused kernel bypass issue. Might be worth comparing approaches if you're still debugging.

[dengls24]

Hi @MELANCHOLY828, thanks for the great catch!

You're right — the KV quantization was indeed not taking effect. The root cause is in kv_cache_quant.py, specifically the condition:

if isinstance(outputs, tuple) and len(outputs) >= 3:

In newer versions of HuggingFace Transformers, when output_attentions=False (the default), the attention module's forward returns a length-2 tuple (attn_output, past_key_value) rather than the length-3 tuple (attn_output, attn_weights, past_key_value). This causes the quantization path to be silently skipped.

Fix (committed in 594f6ea):

1. kv_cache_quant.py — Instead of checking for a fixed tuple length/index, we now scan the output tuple for any object with key_cache and value_cache attributes (i.e., the DynamicCache object). This works regardless of whether the tuple has 2 or 3 elements.
2. exp_niah.py — Added attn_implementation='eager' when loading the quantized model, to prevent fused attention kernels (Flash Attention / cuDNN SDPA) from bypassing the Python-level hooks entirely — as @jagmarques correctly pointed out.

Please pull the latest main and try again. The quantization should now take effect. You can verify by checking that the model outputs differ between the full-precision and quantized runs.

Thanks again for reporting this!