[Bug fix]Fix append attention & noaux_tc unit test

custom_ops/gpu_ops/moe/ep_moe_prefill_func.cu

@@ -28,6 +28,11 @@
 
 #define DISPATCH_NUM_EXPERTS_PER_RANK(num_experts_per_rank, NUM_EXPERTS_PER_RANK, ...)  \
   switch (num_experts_per_rank) {                                                       \
+    case 2: {                                                                           \
+      constexpr size_t NUM_EXPERTS_PER_RANK = 2;                                        \
+      __VA_ARGS__                                                                       \
+      break;                                                                            \
+    }                                                                                   \
     case 8: {                                                                           \
       constexpr size_t NUM_EXPERTS_PER_RANK = 8;                                        \
       __VA_ARGS__                                                                       \

fastdeploy/entrypoints/llm.py

@@ -78,9 +78,9 @@ def __init__(
         chat_template: Optional[str] = None,
         **kwargs,
     ):
+        load_model_register_plugins()
         deprecated_kwargs_warning(**kwargs)
 
-        load_model_register_plugins()
         model = retrive_model_from_server(model, revision)
         tool_parser_plugin = kwargs.get("tool_parser_plugin")
         if tool_parser_plugin:

fastdeploy/model_executor/layers/moe/ep.py

@@ -50,14 +50,30 @@ def get_moe_scores(
     """
     scores = paddle.nn.functional.sigmoid(gating_output)
     scores_with_bias = scores + e_score_correction_bias
-    scores, topk_values, topk_idx = noaux_tc(
-        scores,
-        scores_with_bias,
-        n_group,
-        topk_group,
-        top_k,
-        routed_scaling_factor,
-    )
+    if n_group == 0 or n_group == 1:
+        _, topk_idx = paddle.topk(scores_with_bias, k=top_k, axis=-1) 
+        token_num, top_k = topk_idx.shape
+
+        topk_idx_expanded = paddle.unsqueeze(topk_idx, axis=-1)
+        indices = paddle.concat([
+                paddle.arange(token_num, dtype='int64').unsqueeze(1).tile([1, top_k]).unsqueeze(-1),  # batch_index
+                topk_idx_expanded                                                                     # expert_index
+            ], axis=-1)
+        selected_gate_probs = paddle.gather_nd(scores, indices)
+
+        selected_gate_probs_sum = paddle.sum(selected_gate_probs, axis=1, keepdim=True)
+        topk_weights = selected_gate_probs / selected_gate_probs_sum
+        topk_values = topk_weights * routed_scaling_factor
+        scores = None
+    else:
+        scores, topk_values, topk_idx = noaux_tc(
+            scores,
+            scores_with_bias,
+            n_group,
+            topk_group,
+            top_k,
+            routed_scaling_factor,
+        )
     return scores, topk_values, topk_idx
 
 
@@ -112,6 +128,7 @@ def __init__(
                 low_latency_mode=True,
                 num_qps_per_rank=24,
             )
+            logger.info("Inited deepep_engine")
         # In disaggregated mode on mutiple nodes, we either use
         # high throughput mode or low latency mode.
         else:

fastdeploy/worker/gpu_worker.py

@@ -36,7 +36,8 @@
 
 try:
     ModelRunner = load_model_runner_plugins()
-except:
+except Exception as e:
+    logger.info(f"load_model_runner_plugins encounter error {str(e)}, using default GPUModelRunner")
     from fastdeploy.worker.gpu_model_runner import GPUModelRunner as ModelRunner
 
 

fastdeploy/worker/worker_process.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-
 import argparse
 import json
 import time

tests/layers/test_append_attention.py

@@ -376,7 +376,10 @@ def init_tensor(self):
         )
         self.max_enc_len_this_time = paddle.to_tensor([self.max_enc_len_this_time], "int32", place=paddle.CPUPlace())
         self.max_dec_len_this_time = paddle.to_tensor([self.max_dec_len_this_time], "int32", place=paddle.CPUPlace())
-        self.seq_lens_this_time = self.seq_lens_encoder
+        self.seq_lens_this_time = paddle.to_tensor(
+            self.seq_lens_enc,
+            "int32",
+        )
 
         self.decoder_batch_ids = paddle.full([self.batch_size], 0, dtype="int32")
         self.decoder_tile_ids_per_batch = paddle.full([self.batch_size], 0, dtype="int32")
@@ -567,7 +570,7 @@ def test_all(self):
         )
         # encoder
         # self.seq_lens_encoder,self.seq_lens_decoder,self.max_enc_len_this_time,self.max_dec_len_this_time=get_encoder_decoder_len(self.batch_size,self.seq_len)
-        self.seq_lens_this_time = self.seq_lens_encoder
+        self.seq_lens_this_time[:] = self.seq_lens_encoder[:]
         if self.use_mask_offset:
             print("encoder mask_offset: ", self.mask_offset)
         self.cmp_append_attention(attn_mask=self.attention_mask)

tests/operators/test_noaux_tc.py

@@ -21,12 +21,12 @@ def node_limit_routing(self, gate_probs):
         assert len(gate_probs.shape) == 2
         seq_length, n_experts = gate_probs.shape
 
-        group_scores = gate_probs.reshape([seq_length, 8, -1]).topk(2, axis=-1)[0].sum(axis=-1)
-        group_idx = paddle.topk(group_scores, k=4, axis=-1, sorted=True)[1]
+        group_scores = gate_probs.reshape([seq_length, self.n_group, -1]).topk(2, axis=-1)[0].sum(axis=-1)
+        group_idx = paddle.topk(group_scores, k=self.topk_group, axis=-1, sorted=True)[1]
         group_mask = paddle.zeros_like(group_scores).put_along_axis(
             group_idx, paddle.ones([], dtype="float32"), axis=-1
         )
-        score_mask = group_mask.unsqueeze(-1).expand([seq_length, 8, n_experts // 8]).reshape([seq_length, -1])
+        score_mask = group_mask.unsqueeze(-1).expand([seq_length, self.n_group, n_experts // self.n_group]).reshape([seq_length, -1])
         gate_probs = gate_probs.masked_fill(~score_mask.astype(paddle.bool), float("-inf"))
         return gate_probs
 
@@ -68,8 +68,8 @@ def test_moe_select(self):
 
         ref_topk_values, ref_topk_idx = self.ref_moe_routing()
 
-        paddle.allclose(topk_values, ref_topk_values)
-        paddle.allclose(topk_idx.cast(int), ref_topk_idx.cast(int))
+        assert paddle.allclose(topk_values, ref_topk_values).item()
+        assert paddle.allclose(topk_idx.cast(int), ref_topk_idx.cast(int)).item()
 
 
 if __name__ == "__main__":