[Weight Cache] Enhance test and support EPLB (#36198)

This commit is contained in:
Xun Sun
2026-08-26 23:11:53 -07:00
committed by GitHub
parent b294bd4bc7
commit adcf73d7f7
5 changed files with 77 additions and 9 deletions
-5
View File
@@ -8467,11 +8467,6 @@ class ServerArgs:
"(--weight-cache-mode off) for this configuration."
)
if cfg.weight_cache_mode != "off" and cfg.enable_eplb:
raise ValueError(
"--weight-cache-mode is not supported together with --enable-eplb."
)
def _is_mistral_native_format(self) -> bool:
"""True iff the checkpoint requires load_format=mistral.
+18
View File
@@ -316,6 +316,7 @@ class WeightCacheDaemon:
# The initialized groups are the authority for rank identity. This
# avoids maintaining a second copy of the model-parallel hierarchy.
self._init_distributed(server_args, model_config)
self._initialize_eplb_expert_location_metadata(model_config)
moe_dp_rank = get_parallel().moe_dp_rank
moe_ep_rank = get_parallel().moe_ep_rank
self.config = CacheConfig(
@@ -457,6 +458,23 @@ class WeightCacheDaemon:
f"metadata size ~{total_bytes / 1024 / 1024:.1f} MB"
)
def _initialize_eplb_expert_location_metadata(self, model_config) -> None:
"""Build the same initial physical expert layout as the engine."""
if not self.server_args.enable_eplb:
return
from sglang.srt.eplb.expert_location import (
compute_initial_expert_location_metadata,
set_global_expert_location_metadata,
)
set_global_expert_location_metadata(
compute_initial_expert_location_metadata(
model_config=model_config,
moe_ep_rank=get_parallel().moe_ep_rank,
)
)
def serve(self):
"""Block and serve IPC handles over Unix socket."""
# Do NOT unlink an existing socket here: stale-file cleanup is the launch
@@ -10,6 +10,7 @@ import torch
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
@@ -26,7 +27,7 @@ DEFAULT_MODEL = "Qwen/Qwen3-0.6B"
# IPC handoff is exercised on every PR. Since the CI runner executes the whole
# file per suite, TestWeightCacheDaemonTP2 self-skips when fewer than 2 GPUs are
# visible (i.e. on the 1-gpu runner).
register_cuda_ci(est_time=100, stage="extra-a", runner_config="2-gpu-large")
register_cuda_ci(est_time=280, stage="extra-a", runner_config="2-gpu-large")
register_cuda_ci(est_time=45, stage="base-b", runner_config="1-gpu-small")
# Capture the client server's logs so test_loaded_via_ipc can assert the IPC
@@ -50,9 +51,13 @@ PROMPTS = [
class TestWeightCacheDaemonTP2(CustomTestCase):
"""E2E test: start weight cache daemons, then launch server in client mode with TP2."""
model_override = None
daemon_args = []
server_args = []
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL
cls.model = cls.model_override or DEFAULT_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.tp_size = 2
@@ -74,6 +79,7 @@ class TestWeightCacheDaemonTP2(CustomTestCase):
cls.model,
"--tp-size",
str(cls.tp_size),
*cls.daemon_args,
]
)
@@ -107,6 +113,7 @@ class TestWeightCacheDaemonTP2(CustomTestCase):
str(cls.tp_size),
"--weight-cache-mode",
"client",
*cls.server_args,
],
return_stdout_stderr=(cls.stdout, cls.stderr),
)
@@ -334,5 +341,39 @@ class TestWeightCacheDaemonTP1Smoke(CustomTestCase):
)
class TestWeightCacheDaemonQwen3MoeDP(TestWeightCacheDaemonTP2):
"""Qwen3 with static attention DP through the existing IPC fixture."""
daemon_args = [
"--dp",
"2",
"--ep-size",
"1",
"--enable-dp-attention",
"--enable-dp-lm-head",
"--random-seed",
"42",
]
server_args = daemon_args[:]
class TestWeightCacheDaemonQwen3MoeEP(TestWeightCacheDaemonTP2):
"""Qwen3 MoE with static expert parallelism through the existing IPC fixture."""
model_override = DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN
daemon_args = [
"--dp",
"1",
"--ep-size",
"2",
"--enable-eplb",
"--ep-num-redundant-experts",
"2",
"--random-seed",
"42",
]
server_args = daemon_args[:]
if __name__ == "__main__":
unittest.main()
@@ -269,8 +269,9 @@ class TestDaemonLaunchConfiguration(CustomTestCase):
from sglang.srt.weight_cache import daemon
# The spawn helper receives Engine's already-resolved ServerArgs. A
# minimal namespace keeps this projection test CPU-only and
# model-independent; importantly, no EPLB configuration is involved.
# minimal namespace keeps this forwarding test CPU-only and
# model-independent while covering the static DP/EP and EPLB fields
# consumed by the daemon.
server_args = SimpleNamespace(
model_path="/models/demo",
tp_size=8,
@@ -284,6 +285,8 @@ class TestDaemonLaunchConfiguration(CustomTestCase):
moe_dense_tp_size=1,
moe_a2a_backend="mooncake",
deepep_mode="low_latency",
enable_eplb=True,
ep_num_redundant_experts=72,
load_format="safetensors",
dtype="bfloat16",
quantization="fp8",
@@ -45,6 +45,17 @@ _mock_device.start()
class TestPrepareServerArgs(CustomTestCase):
def test_weight_cache_daemon_allows_static_eplb(self):
args = ServerArgs(
model_path="dummy",
weight_cache_mode="daemon",
enable_eplb=True,
)
# This validation runs before model construction and should allow the
# daemon to build the same static EPLB layout as the engine.
args._handle_load_format()
def test_enable_w4a4_mxfp4_megamoe_sets_deepgemm_env(self):
deepgemm_env = {
"DG_USE_FP4_ACTS": "0",