revert experts

inference/models/mixtral.cc

@@ -167,47 +167,120 @@ void MIXTRAL::create_mixtral_model(FFModel &ff,
     Tensor ff_norm = token_ff_norm[1];
 
     // MoE
-	Tensor w1 = ff.dense(
-        			   ff_norm,
-                       mixtral_config.intermediate_size,
-                       AC_MODE_NONE,
-                       false,
-                       DT_NONE,
-                       nullptr,
-                       nullptr,
-                       nullptr,
-                       REG_MODE_NONE,
-                       0.0f,
-                       std::string("layers." + std::to_string(i) + ".block_sparse_moe_experts_1_w1").c_str());
+    Tensor gate = ff.dense(
+        ff_norm,
+        mixtral_config.num_local_experts,
+        AC_MODE_NONE,
+        false,
+        DT_NONE,
+        nullptr,
+        nullptr,
+        nullptr,
+        REG_MODE_NONE,
+        0.0f,
+        std::string("layers." + std::to_string(i) + ".block_sparse_moe_gate")
+            .c_str());
+    gate = ff.softmax(
+        gate,
+        0,
+        DT_NONE,
+        std::string("layers." + std::to_string(i) + ".block_sparse_moe_softmax")
+            .c_str());
 
-  	Tensor w3 = ff.dense(
-            		   ff_norm,
-                       mixtral_config.intermediate_size,
-                       AC_MODE_NONE,
-                       false,
-                       DT_NONE,
-                       nullptr,
-                       nullptr,
-                       nullptr,
-                       REG_MODE_NONE,
-                       0.0f,
-                       std::string("layers." + std::to_string(i) + ".block_sparse_moe_experts_1_w3").c_str());
+    Tensor topk_out[2] = {nullptr, nullptr};
+    ff.top_k(
+        gate,
+        topk_out,
+        mixtral_config.num_experts_per_tok,
+        false,
+        std::string("layers." + std::to_string(i) + ".block_sparse_moe_topk")
+            .c_str());
+    Tensor topk_values = topk_out[0];
+    Tensor topk_indices = topk_out[1];
 
-  Tensor multi = ff.sigmoid_silu_multi(w1, w3); //DT_NONE,std::string("layers." + std::to_string(i) +".block_sparse_moe_experts." +std::to_string(expert_idx) + "ssm").c_str());
+    Tensor grouped_tokens[mixtral_config.num_local_experts] = {nullptr};
+    ff.group_by(
+        ff_norm,
+        topk_indices,
+        grouped_tokens,
+        mixtral_config.num_local_experts,
+        0.0f,
+        std::string("layers." + std::to_string(i) + ".block_sparse_moe_groupby")
+            .c_str());
 
-  mlp_out = ff.dense(
-      				   multi,
-                       mixtral_config.hidden_size,
-                       AC_MODE_NONE,
-                       false,
-                       DT_NONE,
-                       nullptr,
-                       nullptr,
-                       nullptr,
-                       REG_MODE_NONE,
-                       0.0f,
-                       std::string("layers." + std::to_string(i) + ".block_sparse_moe_experts_1_w2").c_str());
+    Tensor aggregate_inputs[4 + mixtral_config.num_local_experts] = {nullptr};
+    for (int expert_idx = 0; expert_idx < mixtral_config.num_local_experts;
+         expert_idx++) {
+      Tensor w1 = ff.dense(grouped_tokens[expert_idx],
+                           mixtral_config.intermediate_size,
+                           AC_MODE_NONE,
+                           false,
+                           DT_NONE,
+                           nullptr,
+                           nullptr,
+                           nullptr,
+                           REG_MODE_NONE,
+                           0.0f,
+                           std::string("layers." + std::to_string(i) +
+                                       ".block_sparse_moe_experts_" +
+                                       std::to_string(expert_idx) + "_w1")
+                               .c_str());
+
+      Tensor w3 = ff.dense(grouped_tokens[expert_idx],
+                           mixtral_config.intermediate_size,
+                           AC_MODE_NONE,
+                           false,
+                           DT_NONE,
+                           nullptr,
+                           nullptr,
+                           nullptr,
+                           REG_MODE_NONE,
+                           0.0f,
+                           std::string("layers." + std::to_string(i) +
+                                       ".block_sparse_moe_experts_" +
+                                       std::to_string(expert_idx) + "_w3")
+                               .c_str());
+
+      Tensor multi =
+          ff.sigmoid_silu_multi(w1,
+                                w3,
+                                DT_NONE,
+                                std::string("layers." + std::to_string(i) +
+                                            ".block_sparse_moe_experts_" +
+                                            std::to_string(expert_idx) + "ssm")
+                                    .c_str());
+
+      Tensor w2 = ff.dense(multi,
+                           mixtral_config.hidden_size,
+                           AC_MODE_NONE,
+                           false,
+                           DT_NONE,
+                           nullptr,
+                           nullptr,
+                           nullptr,
+                           REG_MODE_NONE,
+                           0.0f,
+                           std::string("layers." + std::to_string(i) +
+                                       ".block_sparse_moe_experts_" +
+                                       std::to_string(expert_idx) + "_w2")
+                               .c_str());
+      aggregate_inputs[4 + expert_idx] = w2;
+    }
+
+    Tensor topk_values_reduced = ff.reduce_sum(topk_values, {0}, true);
+    topk_values = ff.divide(topk_values, topk_values_reduced);
+
+    aggregate_inputs[0] = topk_values;
+    aggregate_inputs[1] = topk_indices;
+    aggregate_inputs[2] = aggregate_inputs[3] = nullptr;
+    mlp_out = ff.aggregate(aggregate_inputs,
+                           mixtral_config.num_local_experts,
+                           0.0f,
+                           std::string("layers." + std::to_string(i) +
+                                       ".block_sparse_moe_experts_aggregate")
+                               .c_str());
   }
+
   // final normalization and linear
   Tensor final_rms_norm_output[2] = {nullptr, nullptr};
   ff.residual_rms_norm(token,