github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 11:32:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Inference/opt] Fused KVCahce Memcopy (#5374)

Browse Source 
* fused kv memcopy

* add TODO in test_kvcache_copy.py
This commit is contained in:
yuehuayingxueluo
2024-02-07 17:15:42 +08:00
committed by GitHub
parent 58740b5f68
commit 6fb4bcbb24
4 changed files with 75 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/inference/modeling/models/nopadding_llama.py
View File

@@ -301,8 +301,9 @@ class NopadLlamaAttention(LlamaAttention):
sm_scale=sm_scale,
)
else:
copy_kv_to_blocked_cache(key_states, k_cache, kv_lengths=sequence_lengths, block_tables=block_tables)
copy_kv_to_blocked_cache(value_states, v_cache, kv_lengths=sequence_lengths, block_tables=block_tables)
copy_kv_to_blocked_cache(
key_states, value_states, k_cache, v_cache, kv_lengths=sequence_lengths, block_tables=block_tables
)
attn_output = flash_decoding_attention(
q=query_states,
k_cache=k_cache,

5
colossalai/inference/modeling/models/padding_llama.py
View File

@@ -356,8 +356,9 @@ class PadLlamaAttention(LlamaAttention):
if attention_mask is not None:
attn_output = pad_input(attn_output, indices, bsz, q_len)
else:
copy_kv_to_blocked_cache(key_states, k_cache, kv_lengths=sequence_lengths, block_tables=block_tables)
copy_kv_to_blocked_cache(value_states, v_cache, kv_lengths=sequence_lengths, block_tables=block_tables)
copy_kv_to_blocked_cache(
key_states, value_states, k_cache, v_cache, kv_lengths=sequence_lengths, block_tables=block_tables
)
attn_output = flash_decoding_attention(
q=query_states,
k_cache=k_cache,