[Question] What are the current KV cache quantization strategies available?

[liye0626]

Hello, I would like to ask what are the current KV cache quantization strategies available? For example, "tensor", "channel", "group", "block", "token", "tensor_group"?

When I set strategy to "tensor", it can run successfully. When set to "group" or "channel", it fails.

• Description in QuantizedKVParameterCache: Quantization strategy (tensor, group, channel) set from Quantization arg's strategy. link: https://github.com/vllm-project/llm-compressor/blob/main/src/llmcompressor/modifiers/quantization/cache.py#L14
• For example: channel
  • code:

kv_cache_dict = {'num_bits': 8, 'type': 'float', 'symmetric': True, 'strategy': 'channel', 'dynamic': False}
recipe = [
    GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"], kv_cache_scheme=kv_cache_dict)
]

-- issue:

  File "/home/user/workspace/pythonlab/llmcompressor/pipelines/sequential/helpers.py", line 72, in forward
    outputs = forward_fn(*args, **kwargs)

  ...

  File "/home/user/workspace/pythonlab/transformers/models/llama/modeling_llama.py", line 309, in forward
    hidden_states, self_attn_weights = self.self_attn(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1857, in _call_impl
    return inner()
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1818, in inner
    hook_result = hook(self, args, result)
  File "/home/user/workspace/pythonlab/llmcompressor/modifiers/utils/hooks.py", line 93, in wrapped_hook
    return hook(*args, **kwargs)
  File "/home/user/workspace/pythonlab/llmcompressor/modifiers/quantization/calibration.py", line 249, in calibrate_kv_cache_output_hook
    update_parameter_data(module, k_scale, KVCacheScaleType.KEY.value)
  File "/home/user/workspace/pythonlab/compressed_tensors/utils/offload.py", line 166, in update_parameter_data
    update_offload_parameter(module, param_name, new_param_data)
  File "/home/user/workspace/pythonlab/compressed_tensors/utils/offload.py", line 257, in update_offload_parameter
    param.data.copy_(data)
RuntimeError: output with shape [1] doesn't match the broadcast shape [1, 1, 1, 1]

[kylesayrs]

Hi there!

Currently only per-tensor KV cache quantization is supported right now, although channel and group/block strategies could be supported on vLLM with some work. https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py#L762-L766

If you want to quantize using these strategies, you can do that using these branches and removing any per_tensor asserts you encounter
#1651
neuralmagic/compressed-tensors#428

Here's an example of per-tensor asymmetric quantization, for example
https://huggingface.co/nm-testing/Meta-Llama-3-8B-Instruct-spinquantR3-FP8_asym-attn

[liye0626]

Hi there!

Currently only per-tensor KV cache quantization is supported right now, although channel and group/block strategies could be supported on vLLM with some work. https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py#L762-L766

If you want to quantize using these strategies, you can do that using these branches and removing any per_tensor asserts you encounter #1651 neuralmagic/compressed-tensors#428

Here's an example of per-tensor asymmetric quantization, for example https://huggingface.co/nm-testing/Meta-Llama-3-8B-Instruct-spinquantR3-FP8_asym-attn

Thanks for your reply. Furthermore, I would like to ask, is it true that the current master branch only supports FP8 KV cache quantization? Does it not support INT8 KV quantization?

[kylesayrs]

@liye0626 Yes, the link above shows that only FP8 per-tensor KV cache quantization is support as of now. However, it is certainly possible to support more KV cache schemes in vLLM, this might be a ideal space for contribution if this is something you're interested in supporting!

[liye0626]

@liye0626 Yes, the link above shows that only FP8 per-tensor KV cache quantization is support as of now. However, it is certainly possible to support more KV cache schemes in vLLM, this might be a ideal space for contribution if this is something you're interested in supporting!

Thanks for your help, it helped me a lot.

[kylesayrs]

Recent support for R3 transforms are also aimed at reducing quantization error for KV cache quantization, making other more extreme quantization schemes more viable.
#1778

[liye0626]

Recent support for R3 transforms are also aimed at reducing quantization error for KV cache quantization, making other more extreme quantization schemes more viable. #1778

Thanks to llm-compressor for its contribution to KV cache quantization, which helped me a lot.

[liye0626]

@kylesayrs Futhermore, I'd like to ask if llm-compressor supports other compression methods (e.g., pruning) besides quantization for KV caches?