chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PD] Support page size > 1 (#5561)

Browse Source
This commit is contained in:
Byron Hsu
2025-04-19 21:54:27 -07:00
committed by GitHub
parent 20f1c8e374
commit ab4b5606e4
4 changed files with 58 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -286,8 +286,12 @@ class MHATokenToKVPool(KVCache):
self.get_key_buffer(i).nbytes for i in range(self.layer_num)
] + [self.get_value_buffer(i).nbytes for i in range(self.layer_num)]
kv_item_lens = [
self.get_key_buffer(i)[0].nbytes for i in range(self.layer_num)
] + [self.get_value_buffer(i)[0].nbytes for i in range(self.layer_num)]
self.get_key_buffer(i)[0].nbytes * self.page_size
for i in range(self.layer_num)
] + [
self.get_value_buffer(i)[0].nbytes * self.page_size
for i in range(self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
# Todo: different memory layout