[Spec] Support MegaMoE for DSpark under dp attention (#34844)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-08-14 17:20:01 -07:00
committed by GitHub
co-authored by weireweire
parent a5ba081fbb
commit 6eb941a34c
2 changed files with 46 additions and 3 deletions
@@ -5,6 +5,7 @@ from sglang.srt.arg_groups.speculative_hook import (
_handle_dspark,
_target_checkpoint_bundles_dspark_draft,
)
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -84,5 +85,34 @@ class TestDsparkDraftPathDefaulting(CustomTestCase):
)
class TestDsparkDpAttentionMoeA2aGate(CustomTestCase):
"""Gate contract for DSpark + dp attention + MoE a2a backends."""
def _dp_server_args(self, *, moe_a2a_backend: str) -> ServerArgs:
server_args = _make_dspark_server_args(
model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config()
)
server_args.enable_dp_attention = True
server_args.enable_dp_lm_head = True
server_args.dp_size = 2
server_args.tp_size = 2
server_args.moe_a2a_backend = moe_a2a_backend
return server_args
def test_only_megamoe_is_admitted(self):
"""Both sides of the allowlist: megamoe passes, others raise by name."""
with envs.SGLANG_RAGGED_VERIFY_MODE.override("static"):
_handle_dspark(self._dp_server_args(moe_a2a_backend="megamoe"))
for backend in ("deepep", "pplx"):
with self.assertRaisesRegex(ValueError, backend):
_handle_dspark(self._dp_server_args(moe_a2a_backend=backend))
def test_a2a_backend_with_compact_verify_mode_raises(self):
server_args = self._dp_server_args(moe_a2a_backend="megamoe")
with envs.SGLANG_RAGGED_VERIFY_MODE.override("compact"):
with self.assertRaisesRegex(ValueError, "static"):
_handle_dspark(server_args)
if __name__ == "__main__":
unittest.main()