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

[Bugfix] fix npu get kv_item_lens in PD separation when use ASCEND_US… (#15852)

Browse Source 
Co-authored-by: ZhengdQin <zhengdqin@gmail.com>
This commit is contained in:
XDaoHong
2026-03-23 15:56:47 +08:00
committed by GitHub
parent 59cb9a9da6
commit 2d288ba8c9
1 changed files with 16 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
View File

@@ -100,13 +100,22 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
self.get_value_buffer(i).nbytes self.get_value_buffer(i).nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num) for i in range(self.start_layer, self.start_layer + self.layer_num)
] ]
kv_item_lens = [ if self.use_fia:
self.get_key_buffer(i)[0].nbytes kv_item_lens = [
for i in range(self.start_layer, self.start_layer + self.layer_num) self.get_key_buffer(i)[0].nbytes * self.page_size
] + [ for i in range(self.start_layer, self.start_layer + self.layer_num)
self.get_value_buffer(i)[0].nbytes ] + [
for i in range(self.start_layer, self.start_layer + self.layer_num) self.get_value_buffer(i)[0].nbytes * self.page_size
] for i in range(self.start_layer, self.start_layer + self.layer_num)
]
else:
kv_item_lens = [
self.get_key_buffer(i)[0].nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i)[0].nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens return kv_data_ptrs, kv_data_lens, kv_item_lens
def set_kv_buffer( def set_kv_buffer(