[DSV4] Support megamoe for CP (#29569)

This commit is contained in:
Siyuan Chen
2026-07-23 14:46:32 -07:00
committed by GitHub
parent 378aea1385
commit 71fe41b6b3
5 changed files with 167 additions and 7 deletions
@@ -37,8 +37,14 @@ _DEEPEP_ENV = {
"SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH": "1",
}
_MEGAMOE_ENV = {
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "8320",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND": "1",
}
class TestDSV4FlashFP4B200Balanced_CP(
class TestDSV4FlashFP4B200Balanced_CP_DeepEP(
SpecDecodingMixin,
BasicDecodeCorrectnessMixin,
GSM8KMixin,
@@ -92,6 +98,55 @@ class TestDSV4FlashFP4B200Balanced_CP(
kill_process_tree(cls.process.pid)
class TestDSV4FlashFP4B200Balanced_CP_Megamoe(
BasicDecodeCorrectnessMixin,
GSM8KMixin,
CustomTestCase,
):
"""Balanced recipe: TP=4, DP=4, DeepEP, EAGLE (1-step spec)."""
gsm8k_accuracy_thres = 0.93
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL)
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--attn-cp-size",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"megamoe",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--deepep-config",
DEEPEP_CONFIG,
],
env=_MEGAMOE_ENV,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
class TestDSV4FlashFP4B200Balanced_CP_NonDeepEP(
SpecDecodingMixin,
BasicDecodeCorrectnessMixin,