github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-12 13:21:47 +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

2
colossalai/inference/core/request_handler.py
View File

@ -227,4 +227,4 @@ class RequestHandler:
self.done_list.extend(finish_seqs) self.done_list.extend(finish_seqs)
return finish_seqs return finish_seqs

5
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),
@ -317,4 +314,4 @@ class PagedAttention:
): ):
return self.pad_decoding_forward( return self.pad_decoding_forward(
q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1), k_cache, v_cache, lengths, block_tables q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1), k_cache, v_cache, lengths, block_tables
) )

2
tests/test_infer/test_inference_engine.py
View File

@ -81,4 +81,4 @@ def test_inference_engine():
if __name__ == "__main__": if __name__ == "__main__":
test_inference_engine() test_inference_engine()