From c9c1d29cf92e05cc2b93bdfb36a4fed772404105 Mon Sep 17 00:00:00 2001 From: zhangdanyang Date: Fri, 14 Mar 2025 11:30:47 +0800 Subject: [PATCH] create parameter(mint.empty) for kvcache --- .../core/parallel_paged_attention_mgr.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/mindformers/experimental/infer/core/parallel_paged_attention_mgr.py b/mindformers/experimental/infer/core/parallel_paged_attention_mgr.py index 85aeb1a29..37ca16488 100644 --- a/mindformers/experimental/infer/core/parallel_paged_attention_mgr.py +++ b/mindformers/experimental/infer/core/parallel_paged_attention_mgr.py @@ -17,9 +17,9 @@ import math import mindspore.common.dtype as mstype -from mindspore import nn, Parameter from mindspore import ops as P -from mindspore.common.initializer import initializer +from mindspore import nn +from mindformers.experimental.infer.core.utils import create_empty_parameter class ParallelPagedAttentionMgr(nn.Cell): @@ -45,10 +45,20 @@ class ParallelPagedAttentionMgr(nn.Cell): self.value_cache = None self.npu_mem_size = npu_mem_size if self.npu_mem_size > 0: - self.key_cache = Parameter(initializer('zeros', kv_shape, compute_dtype), name="key_cache", - requires_grad=False) - self.value_cache = Parameter(initializer('zeros', kv_shape, compute_dtype), name="value_cache", - requires_grad=False) + self.key_cache = create_empty_parameter( + shape=kv_shape, + dtype=compute_dtype, + device="CPU", + name="key_cache", + requires_grad=False, + ) + self.value_cache = create_empty_parameter( + shape=kv_shape, + dtype=compute_dtype, + device="CPU", + name="value_cache", + requires_grad=False, + ) self.reshape_and_cache = P.auto_generate.ReshapeAndCache() self.paged_attention = P.auto_generate.PagedAttention(self.n_heads, self.scale_value, -- Gitee