github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[git] fixed rebased files

pull/5258/head
FrankLeeeee 2024-01-11 13:50:45 +00:00
parent 1513f20f4d
commit 1ded7e81ef
3 changed files with 3 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/inference/modeling/layers/attention.py
View File

@ -58,9 +58,6 @@ def convert_kvcache(cache, lengths, block_tables, pad_id=0):
seq_len = max(lengths)
padded_cache = []
for i in range(bsz):
cache1 = cache[block_tables[i][: needed_blocks[i] - 1]].permute((0, 3, 1, 2)).reshape(-1, num_heads, head_size)
cache2 = cache[block_tables[i][needed_blocks[i] - 1], :, :, : num_remaing_tokens[i]].permute(2, 0, 1)
_cache = torch.cat(
(
cache[block_tables[i][: needed_blocks[i] - 1]].permute((0, 3, 1, 2)).reshape(-1, num_heads, head_size),