kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-19 02:04:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix auto merge & add back get_flat_data_by_layer (#4393)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-03-13 08:46:25 -07:00
committed by GitHub
parent 8e66fbecee
commit a5a892ffd3
3 changed files with 19 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -591,6 +591,9 @@ class MHATokenToKVPoolHost:
def get_flat_data(self, indices):
return self.kv_buffer[:, :, indices]
def get_flat_data_by_layer(self, indices, layer_id):
return self.kv_buffer[:, layer_id, indices]
def assign_flat_data(self, indices, flat_data):
self.kv_buffer[:, :, indices] = flat_data