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