bugfix revise interface get cpu copy for npu mem pool to align with gpu (#29146)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user