multithread multihead.

b4rtaz · Apr 28, 2024 · 9dbfa9f · 9dbfa9f
1 parent 93f4beb
commit 9dbfa9f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/llama2-tasks.cpp b/src/llama2-tasks.cpp
@@ -110,7 +110,7 @@ void llamaMultiheadAttJoin(TASK_ARGS) {
     int nHeadsPerThread = spec->nHeads / nThreads;
 
     int hStart = threadIndex * nHeadsPerThread;
-    int hEnd = threadIndex == nThreads - 1 ? spec->nHeads : hEnd + nHeadsPerThread;
+    int hEnd = (threadIndex == nThreads - 1) ? spec->nHeads : (hStart + nHeadsPerThread);
 
     for (int h = hStart; h < hEnd; h++) {
         // get the query vector for this head