From f458cc5ba47d9e368599137dbfb22168a3fe66c4 Mon Sep 17 00:00:00 2001
From: lzy <569782149@qq.com>
Date: Wed, 3 Dec 2025 16:42:10 +0800
Subject: [PATCH] [Optimization]1.fix tp+ep moe_forward; 2.set
 max_prefill_batch=env.MAX_PREFILL_NUM (#5353)

* [Optimization] 1.fix tp+ep moe_forward; 2.set max_prefill_batch=env.MAX_PREFILL_NUM

* fix test_chunked_moe

---------

Co-authored-by: YuBaoku <49938469+EmmonsCurse@users.noreply.github.com>
---
 tests/distributed/chunked_moe.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/distributed/chunked_moe.py b/tests/distributed/chunked_moe.py
index 86adce5b5..429a9d6c8 100644
--- a/tests/distributed/chunked_moe.py
+++ b/tests/distributed/chunked_moe.py
@@ -138,6 +138,7 @@ class TestChunkedMoE(unittest.TestCase):
         fused_moe = FusedMoE.__new__(FusedMoE)
         fused_moe.ep_size = 2
         fused_moe.tp_size = 1
+        fused_moe.attn_tp_size = 1
         fused_moe.reduce_results = True
 
         fused_moe.fd_config = mock_fd_config