修复不同device之间使用CopyFrom会出错的bug

ztxz16 · Oct 15, 2024 · f6c25e0 · f6c25e0
1 parent 79e40f5
commit f6c25e0
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/src/fastllm.cpp b/src/fastllm.cpp
@@ -304,6 +304,7 @@ namespace fastllm {
     }
 
     void Data::CopyFrom(const Data &ori) {
+        this->ToDevice(ori.dataDevice);
         this->name = ori.name;
         this->isKVCache = ori.isKVCache;
         this->cacheUid = ori.cacheUid;

diff --git a/src/models/minicpm3.cpp b/src/models/minicpm3.cpp
@@ -151,7 +151,6 @@ namespace fastllm {
             Cat(q_nope, q_rope, -1, query_states);
 
             k_rope.Reshape({bsz, seqlen * qk_rope_head_dim});
-            k_rope_expand.ToDevice(DataDevice::CUDA);
             k_rope_expand.CopyFrom(k_rope);
             k_rope_expand.Expansion({bsz, num_attention_heads * seqlen * qk_rope_head_dim});
             for (int i = 1; i < num_attention_heads; i++)
@@ -353,7 +352,6 @@ namespace fastllm {
                 Cat(q_nope, q_rope, -1, query_states);
 
                 k_rope.Reshape({bsz, seqlen * qk_rope_head_dim});
-                k_rope_expand.ToDevice(DataDevice::CUDA);
                 k_rope_expand.CopyFrom(k_rope);
                 k_rope_expand.Expansion({bsz, num_attention_heads * seqlen * qk_rope_head_dim});
                 for (int i = 1; i < num_attention_heads; i++)