[Performance] Use optimized kernels for MQA/GQA

[WoosukKwon]

In theory, MQA/GQA can reduce memory bandwidth for reading KV cache and enable using TensorCore for the dot products in attention mechanism. However, this benefit can be only realized when using optimized kernels that vLLM does not have at the moment.

1. For prefill, vLLM explicitly expands the incoming keys and values before running the attention op:

   vllm/vllm/model_executor/layers/attention.py

   Lines 121 to 128 in e5452dd

   	key = key[:, :,
   	None, :].expand(key.shape[0], self.num_kv_heads,
   	self.num_queries_per_kv,
   	key.shape[-1])
   	value = value[:, :, None, :].expand(value.shape[0],
   	self.num_kv_heads,
   	self.num_queries_per_kv,
   	value.shape[-1])

   because xformers (nor PyTorch SDPA) does not support MQA/GQA at the moment. This is bad for performance since 1) it causes extra overhead of expanding the tensor, and 2) the attention kernel cannot leverage the advantage described above. While FlashAttention efficiently supports MQA/GQA, we need to use it carefully since it does not cover all GPUs/data types/head sizes that xformers supports.
2. For decode, vLLM's current paged attention kernel also does not leverage the benefits of MQA/GQA. To enjoy the benefit, we need to either significantly rewrite the paged attention kernel, or modify the FlashAttention kernel to support paged KV cache.

[casper-hansen]

This is a feature that is high on my list for performance reasons. I have scavenged other people's benchmarks and found an interesting one from TensorRT-LLM that also uses PagedAttention:

Llama 2 7B with 1x A100 80GB:

• num requests: 128
• sequence length: 1024
• input length: 32
• output length: 992

Latency benchmark:

• MHA: 133.40s
• GQA 44.74s
• MQA 29.17s

Conclusion: Due to MHA taking up KV cache, latency increases as you must use more memory bandwidth. This is especially seen in scenarios where we aim for high throughput, i.e. when handling many requests at once. This means we can see a significant improvement in throughput if we spend less memory on the KV cache because of the grouping in GQA.

Reference: NVIDIA/TensorRT-LLM#404

[whitelok]

Seems it is prosible to migrate https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h to fullfill vLLM's requirement.

[WoosukKwon]

@whitelok Thanks for sharing. Unfortunately, it seems the core part of the kernel is provided as cubin files, which are not human-readable.

[whitelok]

sorry. my bad .
you can refer to branch release/0.5.0, it has last opensource code: https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttention112_half.cu

[WoosukKwon]

@whitelok Oh, not really. While I'm not sure at the moment, it seems there is some code that we can leverage. I will look into it soon. Thanks again for sharing!

[whitelok]

seems entrance point is here: https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h#L1015 all vllm need is just modify kvcache buffer.

[beginlner]

FYI, I have modify the FlashAttention kernel to support paged KV cache with a restriction that block_size must match kBlockN in the kernel.

[casper-hansen]

seems entrance point is here: https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h#L1015 all vllm need is just modify kvcache buffer.

It seems they made some changes to GQA/MQA in the new version 0.6.1.

Diff: https://www.diffchecker.com/lixZr3Aj/
New version: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.6.1/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h#L1251

[zhaoyang-star]

have modify the FlashAttention kernel to support paged KV cache with a restriction that block_size must match kBlockN in the kernel.

Great work! As it is very important for performance, could you plan to submit the feature? I am glad to test it.

[casper-hansen]

@zhaoyang-star See the following:

Dao-AILab/flash-attention#678
#1679

[beginlner]

have modify the FlashAttention kernel to support paged KV cache with a restriction that block_size must match kBlockN in the kernel.

Great work! As it is very important for performance, could you plan to submit the feature? I am glad to test it.

I've tidied up the code a bit, you could test using the following two branches:
https://github.com/beginlner/flash-attention/tree/blocked_kvcache
and
https://github.com/beginlner/vllm/tree/blocked_flash_attn

[zhaoyang-star]

have modify the FlashAttention kernel to support paged KV cache with a restriction that block_size must match kBlockN in the kernel.

Great work! As it is very important for performance, could you plan to submit the feature? I am glad to test it.

I've tidied up the code a bit, you could test using the following two branches: https://github.com/beginlner/flash-attention/tree/blocked_kvcache and https://github.com/beginlner/vllm/tree/blocked_flash_attn

Thanks for the info. I will share the latency comparison after benchmark.