From 830d6647c8f080ceaed8392894f90e621fc7e28d Mon Sep 17 00:00:00 2001 From: wufeisheng Date: Sun, 24 Aug 2025 10:58:03 +0800 Subject: [PATCH] fix append attention & noaux_tc unit test --- custom_ops/gpu_ops/moe/ep_moe_prefill_func.cu | 5 +++ fastdeploy/entrypoints/llm.py | 2 +- fastdeploy/model_executor/layers/moe/ep.py | 33 ++++++++++++++----- fastdeploy/worker/gpu_worker.py | 3 +- fastdeploy/worker/worker_process.py | 3 +- test/layers/test_append_attention.py | 7 ++-- test/operators/test_noaux_tc.py | 10 +++--- 7 files changed, 45 insertions(+), 18 deletions(-) diff --git a/custom_ops/gpu_ops/moe/ep_moe_prefill_func.cu b/custom_ops/gpu_ops/moe/ep_moe_prefill_func.cu index d677b360c81..8ae09fd30c1 100644 --- a/custom_ops/gpu_ops/moe/ep_moe_prefill_func.cu +++ b/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__ \ diff --git a/fastdeploy/entrypoints/llm.py b/fastdeploy/entrypoints/llm.py index 001cfad3e00..16c8a35c37d 100644 --- a/fastdeploy/entrypoints/llm.py +++ b/fastdeploy/entrypoints/llm.py @@ -76,9 +76,9 @@ def __init__( enable_logprob: Optional[bool] = False, **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: diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index 02ccead7fb0..d6e0304fbe3 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/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: diff --git a/fastdeploy/worker/gpu_worker.py b/fastdeploy/worker/gpu_worker.py index bfdc92f1dca..9419c2956f7 100644 --- a/fastdeploy/worker/gpu_worker.py +++ b/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 diff --git a/fastdeploy/worker/worker_process.py b/fastdeploy/worker/worker_process.py index 956c89a6662..af6ae743244 100644 --- a/fastdeploy/worker/worker_process.py +++ b/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 @@ -250,6 +249,8 @@ def event_loop_ep(self) -> None: while True: self.worker_healthy_live_signal.value[self.local_rank % self.max_chips_per_node] = int(time.time()) + num_running_requests = 0 + if self.fd_config.parallel_config.tensor_parallel_rank == 0 and self.task_queue.num_tasks() > 0: tasks, read_finish = self.task_queue.get_tasks() diff --git a/test/layers/test_append_attention.py b/test/layers/test_append_attention.py index 1c2ac0bbf62..036824b42f7 100644 --- a/test/layers/test_append_attention.py +++ b/test/layers/test_append_attention.py @@ -375,7 +375,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") @@ -559,7 +562,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[:] self.cmp_append_attention(attn_mask=self.attention_mask) naive_cache_k, naive_cache_v = block_cache_to_naive_cache( self.cache_k, diff --git a/test/operators/test_noaux_tc.py b/test/operators/test_noaux_tc.py index 06e06567379..5442c208bbb 100644 --- a/test/operators/test_noaux_tc.py +++ b/test/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__":