实现graphllm的异步多线程

ztxz16 · Jul 2, 2024 · 5c9645a · 5c9645a
1 parent 83b551f
commit 5c9645a
Show file tree

Hide file tree

Showing 5 changed files with 385 additions and 60 deletions.
diff --git a/include/fastllm.h b/include/fastllm.h
@@ -502,6 +502,8 @@ namespace fastllm {
         Data &operator [] (const std::string &key);
     };
 
+    void *GetExecutor();
+
     void ClearProfiler();
 
     void PrintProfiler();

diff --git a/include/graph.h b/include/graph.h
@@ -58,21 +58,25 @@ namespace fastllm {
         void FusedAttention(ComputeGraphNode &q, ComputeGraphNode &k, ComputeGraphNode &v, 
                             ComputeGraphNode &curk, ComputeGraphNode &curv, 
                             ComputeGraphNode &original, ComputeGraphNode &mask, ComputeGraphNode &output, 
+                            ComputeGraphNode &seqLens,
                             float scale, int maskType, int unitLen); // 融合的attention
         void Linear(ComputeGraphNode &input, ComputeGraphNode &weight, ComputeGraphNode &bias, ComputeGraphNode &output);
         void LlamaRotatePosition2D(ComputeGraphNode &input, ComputeGraphNode &positionIds, ComputeGraphNode &sinData, ComputeGraphNode &cosData, int rotaryDim); // 2D position for llama
         void MulTo(ComputeGraphNode &input0, ComputeGraphNode &input1); // input0 *= input1
         void RMSNorm(ComputeGraphNode &input, ComputeGraphNode &weight, float eps, ComputeGraphNode &output);
         void Silu(ComputeGraphNode &input, ComputeGraphNode &output);
         void Split(ComputeGraphNode &input, int axis, int start, int end, ComputeGraphNode &output);
-        void SplitLastTokenStates(ComputeGraphNode &input, ComputeGraphNode &output);
+        void SplitLastTokenStates(ComputeGraphNode &input, ComputeGraphNode &seqLens, ComputeGraphNode &output);
         void Swiglu(ComputeGraphNode &input, ComputeGraphNode &output);
 
         // 以下op用于调试
         void Exit(); // 退出
         void Print(ComputeGraphNode &input); // 打印
     };
 
+    // 优化计算图
+    void OptimizeComputeGraph(ComputeGraph &graph, WeightMap &weight);
+
     // 执行计算图
     void RunComputeGraph (const ComputeGraph &graph, 
                         const std::map <std::string, int> &deviceMap,

diff --git a/src/fastllm.cpp b/src/fastllm.cpp
@@ -2903,6 +2903,10 @@ namespace fastllm {
         }
     }
 
+    void *GetExecutor() {
+        return (void*)curExecutor;
+    }
+
     void ClearProfiler() {
         curExecutor->ClearProfiler();
     }