github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-11 21:01:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[git] fixed rebased files

Browse Source
This commit is contained in:
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) seq_len = max(lengths)
padded_cache = [] padded_cache = []
for i in range(bsz): 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 = torch.cat(
( (
cache[block_tables[i][: needed_blocks[i] - 1]].permute((0, 3, 1, 2)).reshape(-1, num_heads, head_size), cache[block_tables[i][: needed_blocks[i] - 1]].permute((0, 3, 1, 2)).reshape(-1, num_heads, head_size),