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

bugfix revise interface get cpu copy for npu mem pool to align with gpu (#27802)

Browse Source
This commit is contained in:
McZyWu
2026-06-15 17:24:10 +08:00
committed by GitHub
parent 1180b70440
commit bf186cf8fc
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+2 -2
View File
@@ -245,7 +245,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# NPUMHATokenToKVPool stores buffers as # NPUMHATokenToKVPool stores buffers as
# (num_pages, page_size, head_num, head_dim) # use_fia=False # (num_pages, page_size, head_num, head_dim) # use_fia=False
# (num_pages*page_size, 1, head_num, head_dim) # use_fia=True # (num_pages*page_size, 1, head_num, head_dim) # use_fia=True
def get_cpu_copy(self, indices): def get_cpu_copy(self, indices, mamba_indices=None):
torch.npu.synchronize() torch.npu.synchronize()
buf_of_layers = [] buf_of_layers = []
for local_layer_id in range(self.layer_num): for local_layer_id in range(self.layer_num):
@@ -260,7 +260,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
torch.npu.synchronize() torch.npu.synchronize()
return kv_cache_cpu return kv_cache_cpu
def load_cpu_copy(self, kv_cache_cpu, indices): def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
torch.npu.synchronize() torch.npu.synchronize()
chunk_size = self.cpu_offloading_chunk_size chunk_size = self.cpu_offloading_chunk_size
for local_layer_id in range(self.layer_num): for local_layer_id in range(self.layer_num):