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

add paged-attetionv2: support seq length split across thread block (#5707)

Browse Source
This commit is contained in:
Steve Luo
2024-05-14 12:46:54 +08:00
committed by GitHub
parent 18d67d0e8e
commit 7806842f2d
8 changed files with 704 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/inference/modeling/models/nopadding_baichuan.py
View File

@@ -338,7 +338,8 @@ class NopadBaichuanAttention(ParallelModule):
block_size,
kv_seq_len,
fd_inter_tensor.mid_output,
fd_inter_tensor.mid_output_lse,
fd_inter_tensor.exp_sums,
fd_inter_tensor.max_logits,
self.alibi_slopes,
sm_scale,
)

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

@@ -596,7 +596,8 @@ class NopadLlamaAttention(LlamaAttention, ParallelModule):
block_size,
kv_seq_len,
fd_inter_tensor.mid_output,
fd_inter_tensor.mid_output_lse,
fd_inter_tensor.exp_sums,
fd_inter_tensor.max_logits,
None,
sm_scale,
)