tfacc attention speedup

ztxz16 · May 8, 2024 · cf1f6b1 · cf1f6b1
1 parent 0d9765a
commit cf1f6b1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/devices/tfacc/fastllm-tfacc.cpp b/src/devices/tfacc/fastllm-tfacc.cpp
@@ -368,12 +368,13 @@ namespace fastllm {
         U8Buffer buffer;
         buffer.WriteInt(configString.size());
         buffer.WriteBytes((uint8_t*)configString.data(), configString.size());
-        buffer.WriteBytes(q->cpuData, q->GetBytes());
+        //buffer.WriteBytes(q->cpuData, q->GetBytes());
 
-        memcpy((uint8_t*)this->buf, buffer.buffer.data(), buffer.buffer.size());
+        RunMultiThreadMemcpy((uint8_t*)this->buf, buffer.buffer.data(), buffer.buffer.size(), GetAlivePool());
+        RunMultiThreadMemcpy((uint8_t*)this->buf + buffer.buffer.size(), q->cpuData, q->GetBytes(), GetAlivePool());
         this->Launch(opType);
         this->Wait();
 
-        memcpy(output->cpuData, (uint8_t*)result, output->GetBytes());
+        RunMultiThreadMemcpy(output->cpuData, (uint8_t*)result, output->GetBytes(), GetAlivePool());
     }
 }
diff --git a/third_party/tfacc/server b/third_party/tfacc/server