Adding PagedAttention support for CausalLM models - #982
Conversation
| v_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out) | ||
| return k_out, v_out | ||
|
|
||
| def read_only_pagedAttention(self, block_index, updated, cache_kwargs): |
There was a problem hiding this comment.
nit: can we rename this to read_only_paged_attention()?
There was a problem hiding this comment.
I was staying consistent with the actual name of the technique from the paper: https://arxiv.org/pdf/2309.06180
The attention mechanism is called PagedAttention rather than paged attention, hence I was keeping pagedAttention in our naming. I can change to paged_attention for all the methods if that would look better with snake case convention.
| v_out = torch.where((invalid_mask.unsqueeze(1)).unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), v_out) | ||
| return k_out, v_out | ||
|
|
||
| def write_only_pagedAttention(self, key_states, value_states, cache_kwargs): |
There was a problem hiding this comment.
write_only_paged_attention()?
| """ | ||
| return self.layers[layer_idx].read_only_blockedKV(start_index, end_index, cache_kwargs) | ||
|
|
||
| def read_only_pagedAttention(self, block_index, updated, layer_idx, cache_kwargs): |
|
nit: lint and format are missing, pls check. |
5dc796a to
ee2dadb
Compare
|
@vaibverm lint, format and unit tests seems to be failing again, pls check, rest LGTM :) thanks |
|
@vbaddi - Pushed changes for lint/format fixes. Should resolve unit test fail as well (was a syntax issue). Running unit tests locally as well to confirm fixes. |
|
@quic-rishinr, @ochougul - Would love to get review feedback from you as well. Please let me know if I can help with clarifying anything. |
8273e7f to
e9330a8
Compare
Issue: Generated text is not accurate (unable to identify object in given image) 5/25: Found a workaround, seems like compiler issue - debugging further --------- Signed-off-by: vtirumal <vtirumal@qti.qualcomm.com> Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
changed the code from doing the exact same math repeatedly. Signed-off-by: Anuj Gupta <anujgupt@qti.qualcomm.com> Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
…ormat cleanup Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
Signed-off-by: Vaibhav Verma <vaibverm@qti.qualcomm.com>
|
Cloned at: #1209 |
5 participants
This PR adds the PagedAttention (https://arxiv.org/pdf/2309.06180) support for all CausalLM models in QEfficient.
The major change is that KV cache is not treated as a contiguous memory under this implementation but rather a collection of blocks which can reside in a non-contiguous fashion inside the memory. This forces cache scatter and gather operations to happen per KV block.
Summary of changes compared to BlockedKV:
4) a) block_id is each entry in the block_table and points to the physical K/V block that needs to be read/written corresponding to (position_id // kv_block_size)th entry in block_table. ‘-1’ signifies invalid/unallocated block.
4) b) slot_id tells how many entries are already filled in currently active block => read up to / write after (slot_id – 1)