From bb7d3440b56b3a054bfa34aedd100a40d4451cd9 Mon Sep 17 00:00:00 2001 From: McZyWu Date: Mon, 29 Jun 2026 19:30:14 +0800 Subject: [PATCH] bugfix revise interface get cpu copy for npu mem pool to align with gpu (#29146) --- 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 7e527c6f0..f60083929 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -506,7 +506,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): out.append(layer_chunks) return out - def get_cpu_copy(self, indices): + def get_cpu_copy(self, indices, mamba_indices=None): torch.npu.synchronize() buf_of_layers = [] has_ik = self.index_head_dim is not None @@ -524,7 +524,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): 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 has_ik = self.index_head_dim is not None