revert rope parallel

intel · Apr 30, 2024 · 7dc4dd8 · 7dc4dd8
1 parent a861d4e
commit 7dc4dd8
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/neural_speed/core/ne_layers.c b/neural_speed/core/ne_layers.c
@@ -11493,11 +11493,10 @@ void ne_graph_compute(struct ne_context* ctx, struct ne_cgraph* cgraph) {
         case NE_OP_DIAG_MASK_INF:
         case NE_OP_PADDING_MASK_INF:
         case NE_OP_ROPE:
-          // only first token use parallel
-          if (node->type == NE_TYPE_BTLA)
-            node->n_tasks = 1;
-          else
+          if (node->type != NE_TYPE_BTLA && node->src0->ne[2] > 1)
             node->n_tasks = n_threads;
+          else
+            node->n_tasks = 1;
           break;
         case NE_OP_SOFT_MAX: {
           size_t rows = ne_nrows(node->src0);