From bf186cf8fc05b687c9f00e47f283760d2996b244 Mon Sep 17 00:00:00 2001 From: McZyWu Date: Mon, 15 Jun 2026 17:24:10 +0800 Subject: [PATCH] bugfix revise interface get cpu copy for npu mem pool to align with gpu (#27802) --- python/sglang/srt/hardware_backend/npu/memory_pool_npu.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index c3b9dfd64..c81d0f819 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -245,7 +245,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): # NPUMHATokenToKVPool stores buffers as # (num_pages, page_size, head_num, head_dim) # use_fia=False # (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() buf_of_layers = [] for local_layer_id in range(self.layer_num): @@ -260,7 +260,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): torch.npu.synchronize() 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() chunk_size = self.cpu_offloading_chunk_size for local_layer_id in range(self.layer_num):