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

[PD] Fix large page size + chunk prefill (#5588)

Browse Source
This commit is contained in:
Byron Hsu
2025-04-20 17:21:54 -07:00
committed by GitHub
parent dcb8232596
commit c951d312ed
5 changed files with 37 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -446,13 +446,16 @@ class MLATokenToKVPool(KVCache):
]
self.layer_transfer_counter = None
self.page_size = page_size
# for disagg
def get_contiguous_buf_infos(self):
# MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in range(self.layer_num)]
kv_data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
kv_item_lens = [self.kv_buffer[i][0].nbytes for i in range(self.layer_num)]
kv_item_lens = [
self.kv_buffer[i][0].nbytes * self.page_size for i in range(self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def get_key_buffer(self, layer_id: int):