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

adapted to pad_context_forward

Browse Source
This commit is contained in:
yuehuayingxueluo
2024-01-09 13:52:53 +08:00
committed by FrankLeeeee
parent 47e53eaa1c
commit fa4fbdbffb
9 changed files with 42 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/inference/kv_cache/kvcache_manager.py
View File

@@ -112,7 +112,7 @@ class KVCacheManager:
def get_kv_cache(self):
"""Get k_cache and v_cache"""
return self._kv_caches[0], self._kv_caches[1]
return self._kv_caches
def get_max_blocks_per_sequence(self) -> int:
"""Get the maximum number of blocks that can be allocated for a single sequence."""