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

Fix NPUMHATokenToKVPool missing k_data_ptrs/v_data_ptrs (#35820)

Browse Source
This commit is contained in:
ming_wang
2026-08-26 11:01:11 +08:00
committed by GitHub
parent 5bcc978f2a
commit 2e4773aadd
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/pool_host/mha.py
+10 -6
View File
@@ -128,12 +128,16 @@ class MHATokenToKVPoolHost(HostKVCache):
) )
if self.mtp_draft_device_pools: if self.mtp_draft_device_pools:
device_pools = (self.device_pool, *self.mtp_draft_device_pools) device_pools = (self.device_pool, *self.mtp_draft_device_pools)
self.packed_device_k_data_ptrs = torch.cat( if not _is_npu:
[pool.k_data_ptrs for pool in device_pools] self.packed_device_k_data_ptrs = torch.cat(
) [pool.k_data_ptrs for pool in device_pools]
self.packed_device_v_data_ptrs = torch.cat( )
[pool.v_data_ptrs for pool in device_pools] self.packed_device_v_data_ptrs = torch.cat(
) [pool.v_data_ptrs for pool in device_pools]
)
else:
self.packed_device_k_data_ptrs = None
self.packed_device_v_data_ptrs = None
self.packed_device_k_buffers = [ self.packed_device_k_buffers = [
buffer for pool in device_pools for buffer in pool.k_buffer buffer for pool in device_pools for buffer in pool.k_buffer
] ]