Optimize LlamaAttention by fusing QKV projections#40092
Optimize LlamaAttention by fusing QKV projections#40092null-pointer-access wants to merge 1 commit into
Conversation
|
[For maintainers] Suggested jobs to run (before merge) run-slow: llama |
ArthurZucker
left a comment
There was a problem hiding this comment.
Hey!
Thanks but we cannoit accept this.
Specifically, when running a single prefill with seq_len=128 on 1xH100, it achieves 28% efficiency improvement.
Well, this needs benchmarking, making sure it holds with compile, with cudagraphs, with fa2 kernels, etc
|
@ArthurZucker Thanks for the feedback! Actually, this implementation is inspired by which replaces multiple operator calls (q_proj, k_proj, v_proj) with a single fused operator (qkv_proj). To show that this fix can provide stable speedup under different workloads, we did extensive evaluations and here's the result with TinyLlama-1.1B-Chat-v1.0 on an H100 GPU:
You can find the benchmark script here. |
| qkv = self.qkv_proj(hidden_states) | ||
| query_pos = self.config.num_attention_heads * self.head_dim | ||
| query_states = qkv[..., :query_pos] | ||
| key_states = qkv[..., query_pos : query_pos + self.config.num_key_value_heads * self.head_dim] | ||
| value_states = qkv[..., query_pos + self.config.num_key_value_heads * self.head_dim :] |
There was a problem hiding this comment.
Nice work! I agree that fusing the linear layers for Q, K, and V should generally be faster than using three separate ones. However, there's a potential consequence to consider.
This fusion will result in the q, k, v tensors not being contiguous in memory. As a result, the attention mechanism will need to call contiguous() on them afterward, which involves copying the non-contiguous q, k, v tensors to a new contiguous memory location. This copying could introduce a latency and performance trade-off that should be evaluated in benchmarks
There was a problem hiding this comment.
Hi @ducviet00 , thanks for your feedback! I looked into the problem you mentioned and did some benchmark, and find that there would not be additional contiguous() because in LlamaModel, these tensors will then be forced to be contiguous by torch.cat or repeat_kv.
transformers/src/transformers/models/llama/modeling_llama.py
Lines 109 to 113 in 7cbcd9b
transformers/src/transformers/models/llama/modeling_llama.py
Lines 159 to 168 in 7cbcd9b
ArthurZucker
left a comment
There was a problem hiding this comment.
Hey @null-pointer-access 🤗
It's a very well well known "issue" in the sense that it can be faster but as your benchmarks show, it is not obvious!
If you think about inference: cudagraph + compile make it ~same perf
If you think about trining: the fact that you split q, k, v make it a lost simpler to run PEFT (LORA).
And more than that, its impossible to change the format now, because it would break the 300K models on the hub 😉
There is the deprecated/open_llama that tried to push the performances, but TBH it does not get usage.
I am not convinced that it is worth it expecially for all torch dtype / compile / cudagraph / tensor parallel etc!
|
Not to chime into a conversation, but perhaps you could package the changes in the PR into a cookbook-style recipe here: https://github.com/huggingface/cookbook The recipe could document how LlamaAttention or any other module can be speed up and be modified for use in transformers, for a particular use-case. |
What does this PR do?
This PR fuses the q, k, and v projections in LlamaAttention to a single qkv projection. This can improve the efficiency and GPU utilization when the batch size is small. Specifically, when running a single prefill with
seq_len=128on 1xH100, it achieves 28% efficiency improvement.Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@ArthurZucker