diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 2a09a0c8c0ce..784d64fb9464 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1626,7 +1626,7 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec `--speculative-moe-a2a-backend` - MOE A2A backend for EAGLE speculative decoding, see `--moe-a2a-backend` for options. Same as moe a2a backend if unset. + MoE A2A backend for speculative decoding; see `--moe-a2a-backend` for options. If unset or `none`, inherit the target model's backend. Explicit `none` does not select a separate local dispatcher when the target uses another backend. `None` none, deepep, mooncake, nixl, mori, ascend_fuseep, flashinfer, megamoe, pplx diff --git a/docs/docs/advanced_features/speculative_decoding.mdx b/docs/docs/advanced_features/speculative_decoding.mdx index 004e9ef75b81..6d0e3bb04076 100644 --- a/docs/docs/advanced_features/speculative_decoding.mdx +++ b/docs/docs/advanced_features/speculative_decoding.mdx @@ -946,7 +946,7 @@ Below is a comprehensive list of all speculative decoding parameters available i --speculative-moe-a2a-backend str None - MoE all-to-all backend for the draft model + MoE all-to-all backend for the draft model. If unset or none, inherit --moe-a2a-backend; other values select an explicit draft backend. --speculative-draft-model-quantization diff --git a/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx b/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx index d66366ba923c..369fd7229c69 100644 --- a/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx +++ b/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx @@ -1476,7 +1476,7 @@ non-default speculative acceptance thresholds or deterministic inference. `--speculative-moe-a2a-backend` `None` - `ascend_fuseep` (the only supported value on Ascend NPU) + `none` (inherit the target backend), `deepep`, `ascend_fuseep`; model and quantization support varies by backend. A2/A3 Series diff --git a/python/sglang/srt/arg_groups/fields/spec.py b/python/sglang/srt/arg_groups/fields/spec.py index ac91102fdcf5..58178adef437 100644 --- a/python/sglang/srt/arg_groups/fields/spec.py +++ b/python/sglang/srt/arg_groups/fields/spec.py @@ -170,7 +170,10 @@ class Spec(msgspec.Struct): speculative_moe_a2a_backend: A[ Optional[str], Arg( - help="Choose the backend for MoE A2A in speculative decoding", + help=( + "Choose the backend for MoE A2A in speculative decoding. " + "If unset or 'none', inherit --moe-a2a-backend." + ), choices=[ "none", "deepep", diff --git a/python/sglang/srt/arg_groups/moe_hook.py b/python/sglang/srt/arg_groups/moe_hook.py index d8e9ae02df12..f4708d481c19 100644 --- a/python/sglang/srt/arg_groups/moe_hook.py +++ b/python/sglang/srt/arg_groups/moe_hook.py @@ -527,7 +527,7 @@ def validate_deepep_v2_speculative_draft(server_args: Any) -> None: """Reject an explicit or inherited DeepEP v2 draft backend.""" view = resolved_view(server_args) draft_backend = view.speculative_moe_a2a_backend - if draft_backend is None and view.speculative_algorithm: + if draft_backend in (None, "none") and view.speculative_algorithm: from sglang.srt.speculative.spec_info import SpeculativeAlgorithm algorithm = SpeculativeAlgorithm.from_string(view.speculative_algorithm) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 0d33162e0e7a..a48dba494c94 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -608,7 +608,7 @@ def _handle_dspark(server_args: ServerArgs) -> None: ) if ( not _is_npu - and cfg.speculative_moe_a2a_backend is not None + and cfg.speculative_moe_a2a_backend not in (None, "none") and cfg.speculative_moe_a2a_backend != cfg.moe_a2a_backend ): raise ValueError( diff --git a/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py b/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py index 22c1291608c6..fa0f702c3591 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py +++ b/python/sglang/srt/hardware_backend/npu/moe/finalize_routing.py @@ -42,6 +42,10 @@ def _finalize_routing( expanded_row_idx: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: + if self.drop_pad_mode == 3 and hidden_states.ndim == 2: + # Drop mode requires [E, C, H], but uses flattened row indices. + # View the packed GMM rows as one capacity buffer without copying. + hidden_states = hidden_states.unsqueeze(0) return torch.ops.npu.npu_moe_finalize_routing( hidden_states, skip1=None, diff --git a/python/sglang/srt/hardware_backend/npu/moe/init_routing.py b/python/sglang/srt/hardware_backend/npu/moe/init_routing.py index f4de357c6f4f..b8d0222834df 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/init_routing.py +++ b/python/sglang/srt/hardware_backend/npu/moe/init_routing.py @@ -96,6 +96,7 @@ def _init_routing( topk_ids: torch.Tensor, num_experts: int, top_k: int, + active_expert_range: Optional[Tuple[int, int]] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: num_tokens = hidden_states.shape[0] hidden_states, expanded_row_idx, expert_tokens, pertoken_scale = ( @@ -106,7 +107,11 @@ def _init_routing( expert_num=num_experts, expert_tokens_num_type=1, expert_tokens_num_flag=True, - active_expert_range=[0, num_experts], + active_expert_range=( + list(active_expert_range) + if active_expert_range is not None + else [0, num_experts] + ), quant_mode=self.quant_mode, ) ) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py b/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py index b8303a4fabb3..c448889ad7de 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/ascend_tp.py @@ -52,6 +52,14 @@ class AscendTPDispatcher(BaseDispatcher): def __init__(self, moe_runner_config: MoeRunnerConfig): super().__init__() self.num_experts = moe_runner_config.num_experts + self.num_local_experts = moe_runner_config.num_local_experts + self.num_local_shared_experts = moe_runner_config.num_fused_shared_experts + self.num_local_routed_experts = ( + self.num_local_experts - self.num_local_shared_experts + ) + self.moe_ep_size = get_parallel().moe_ep_size + self.moe_ep_rank = get_parallel().moe_ep_rank + self.local_expert_mapping = None self.top_k = moe_runner_config.top_k self._dispatch_output: Optional[AscendTPDispatchOutput] = None @@ -76,18 +84,21 @@ def set_quant_config(self, quant_config: dict) -> None: def set_ascend_dispatcher_output_dtype(self) -> None: """Choose init & finalize routing kernels based on quant config.""" self.ascend_dispatcher_output_dtype = get_ascend_dispatcher_output_dtype(self) + # EP filtering produces -1 row indices for nonlocal experts. Mode 3 + # supports dropped routes in row-major order; mode 2 does not. + drop_pad_mode = 3 if self.moe_ep_size > 1 else 2 if self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.BF16: self.init = NPUMoEInitRouting_v2(quant_mode=-1) - self.finalize = NPUFinalizeRouting(drop_pad_mode=2) + self.finalize = NPUFinalizeRouting(drop_pad_mode=drop_pad_mode) self.group_list_type = 1 elif self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.INT8: self.init = NPUMoEInitRouting_v2(quant_mode=1) - self.finalize = NPUFinalizeRouting(drop_pad_mode=2) + self.finalize = NPUFinalizeRouting(drop_pad_mode=drop_pad_mode) self.group_list_type = 1 elif self.ascend_dispatcher_output_dtype == DispatcherOutputDtype.MXFP8: self.init = NPUMoEInitRouting_v2(quant_mode=MXFP8_QUANT_MODE) - self.finalize = NPUFinalizeRouting(drop_pad_mode=2) + self.finalize = NPUFinalizeRouting(drop_pad_mode=drop_pad_mode) self.group_list_type = 1 else: raise ValueError( @@ -102,6 +113,37 @@ def dispatch( topk_ids = topk_ids.to(torch.int32) top_k = topk_weights.shape[-1] + if self.moe_ep_size > 1: + if self.local_expert_mapping is None: + # Unlike the GPU dispatcher, use a nonnegative drop sentinel: + # init_routing_v2 filters IDs outside active_expert_range. + # The extra entry also maps padding IDs (-1) to that sentinel. + self.local_expert_mapping = torch.full( + (self.num_experts + 1,), + self.num_local_experts, + dtype=torch.int32, + device=topk_ids.device, + ) + start = self.moe_ep_rank * self.num_local_routed_experts + self.local_expert_mapping[ + start : start + self.num_local_routed_experts + ] = torch.arange( + self.num_local_routed_experts, + dtype=torch.int32, + device=topk_ids.device, + ) + if self.num_local_shared_experts > 0: + self.local_expert_mapping[ + self.num_experts + - self.num_local_shared_experts : self.num_experts + ] = torch.arange( + self.num_local_routed_experts, + self.num_local_experts, + dtype=torch.int32, + device=topk_ids.device, + ) + topk_ids = self.local_expert_mapping[topk_ids] + ( permuted_hidden_states, expanded_row_idx, @@ -112,6 +154,7 @@ def dispatch( topk_ids, self.num_experts, top_k, + active_expert_range=(0, self.num_local_experts), ) self._dispatch_output = AscendTPDispatchOutput( diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 9840e671f863..9bf7c25f9d43 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -460,7 +460,7 @@ def initialize_moe_config(): ) moe.speculative_a2a_backend = ( MoeA2ABackend(spec.speculative_moe_a2a_backend) - if spec.speculative_moe_a2a_backend is not None + if spec.speculative_moe_a2a_backend not in (None, "none") else moe.a2a_backend ) moe.deepep_mode = DeepEPMode(exec_moe.deepep_mode) diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 50e1ad2ed8ef..78e8b1545531 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -288,9 +288,9 @@ def refresh_deep_gemm_layout_memory_budget( get_spec().speculative_moe_runner_backend or get_exec().moe.moe_runner_backend ) - moe_a2a_backend = ( - get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend - ) + moe_a2a_backend = get_spec().speculative_moe_a2a_backend + if moe_a2a_backend in (None, "none"): + moe_a2a_backend = get_exec().moe.moe_a2a_backend else: moe_runner_backend = get_exec().moe.moe_runner_backend moe_a2a_backend = get_exec().moe.moe_a2a_backend diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 6ef4acbb32b0..4bc635330e59 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -72,9 +72,9 @@ def should_run_flashinfer_autotune( get_spec().speculative_moe_runner_backend or get_exec().moe.moe_runner_backend ) - a2a_backend_str = ( - get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend - ) + a2a_backend_str = get_spec().speculative_moe_a2a_backend + if a2a_backend_str in (None, "none"): + a2a_backend_str = get_exec().moe.moe_a2a_backend else: backend_str = get_exec().moe.moe_runner_backend a2a_backend_str = get_exec().moe.moe_a2a_backend diff --git a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_none_dispatcher.py b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_none_dispatcher.py new file mode 100644 index 000000000000..5241d9d6a6b7 --- /dev/null +++ b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_none_dispatcher.py @@ -0,0 +1,148 @@ +"""Check local EP routing against a dense reference using real NPU kernels. + +One NPU runs each logical EP rank sequentially; summing its partial outputs +models the all-reduce performed by the model after the ``none`` dispatcher. +""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig +from sglang.srt.layers.moe.token_dispatcher.ascend_tp import ( + AscendTPCombineInput, + AscendTPDispatcher, +) +from sglang.srt.layers.moe.topk import StandardTopKOutput +from sglang.srt.utils import is_npu +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.test_utils import CustomTestCase + +register_npu_ci(est_time=30, suite="full-1-npu-a3", nightly=True) + + +@unittest.skipUnless(is_npu(), "Requires Ascend NPU kernels") +class TestNpuNoneDispatcher(CustomTestCase): + def _check_routing(self, ids, num_routed=4, ep_size=2, num_shared=0, int8=False): + ids = torch.tensor(ids, dtype=torch.int32).reshape(-1, 2) + num_tokens, top_k = ids.shape + num_experts = num_routed + num_shared + local_routed = num_routed // ep_size + num_local = local_routed + num_shared + # Binary fractions make BF16 matmul/combine exact; distinct expert + # transforms detect a wrong rank offset even when count shapes match. + x = ((torch.arange(num_tokens * 32) % 13 - 6) / 8).reshape(num_tokens, 32) + weights = torch.stack( + [torch.eye(32).roll(expert % 32, dims=1) for expert in range(num_experts)] + ) + scores = torch.tensor([0.25, 0.75]).repeat(num_tokens, 1) + expected = torch.zeros_like(x) + for row in range(num_tokens): + for slot in range(top_k): + expert = int(ids[row, slot]) + if expert >= 0: + expected[row] += scores[row, slot] * (x[row] @ weights[expert]) + + total = torch.zeros_like(x) + for rank in range(ep_size): + with self.subTest(rank=rank): + config = MoeRunnerConfig( + num_experts=num_experts, + num_local_experts=num_local, + num_fused_shared_experts=num_shared, + top_k=top_k, + ) + with patch( + "sglang.srt.layers.moe.token_dispatcher.ascend_tp.get_parallel", + return_value=SimpleNamespace(moe_ep_size=ep_size, moe_ep_rank=rank), + ): + dispatcher = AscendTPDispatcher(config) + if int8: + dispatcher.set_quant_config({"dispatcher_output_dtype": "int8"}) + local_ids = list(range(rank * local_routed, (rank + 1) * local_routed)) + local_ids += list(range(num_routed, num_experts)) + local_weights = weights[local_ids].clone() + # Shared experts are replicated across EP ranks. Their caller + # supplies the 1/EP scaling; the dispatcher must keep them local. + if num_shared: + local_weights[-num_shared:] /= ep_size + local_weights = local_weights.to(device="npu", dtype=torch.bfloat16) + topk = StandardTopKOutput( + topk_weights=scores.to(device="npu", dtype=torch.bfloat16), + topk_ids=ids.npu(), + router_logits=None, + ) + hidden = x.to(device="npu", dtype=torch.bfloat16) + # Repeat to cover cached expert mappings and cleared combine state. + for _ in range(2): + dispatched = dispatcher.dispatch(hidden, topk) + counts = dispatched.expert_tokens.cpu() + self.assertEqual(counts.numel(), local_weights.shape[0]) + torch.testing.assert_close( + counts, + torch.tensor([(ids == expert).sum() for expert in local_ids]), + ) + permuted = dispatched.hidden_states + if int8: + self.assertEqual(permuted.dtype, torch.int8) + permuted = ( + permuted.float() + * dispatched.hidden_states_scale.float().reshape(-1, 1) + ).to(torch.bfloat16) + expert_output = torch.ops.npu.npu_grouped_matmul( + [permuted], + [local_weights], + group_list=dispatched.expert_tokens, + group_type=0, + split_item=3, + group_list_type=dispatched.group_list_type, + )[0] + # Unused GMM rows are undefined. Poison them so a finalizer + # that accidentally reads dropped routes cannot pass by luck. + expert_output[int(counts.sum()) :].fill_(float("nan")) + actual = dispatcher.combine(AscendTPCombineInput(expert_output)) + partial = torch.zeros_like(x) + for row in range(num_tokens): + for slot in range(top_k): + expert = int(ids[row, slot]) + if expert in local_ids: + scale = 1 / ep_size if expert >= num_routed else 1 + partial[row] += ( + scores[row, slot] + * scale + * (x[row] @ weights[expert]) + ) + torch.testing.assert_close( + actual.float().cpu(), partial, rtol=0, atol=0.008 if int8 else 0 + ) + total += actual.float().cpu() + torch.testing.assert_close(total, expected, rtol=0, atol=0.008 if int8 else 0) + + def test_ep1(self): + self._check_routing([[0, 2], [1, 3]], ep_size=1) + + def test_ep2(self): + self._check_routing([[0, 2], [1, 3], [2, 3], [0, 1]]) + + def test_ep2_qwen36_expert_count(self): + self._check_routing( + [[0, 127], [128, 255], [255, 1], [129, 126]], num_routed=256 + ) + + def test_ep2_no_local_tokens(self): + self._check_routing([[2, 3]]) + + def test_ep2_shared_experts_and_padding(self): + self._check_routing([[0, 4], [2, 4], [-1, 4]], num_shared=1) + + def test_ep2_int8_routing(self): + self._check_routing([[0, 2], [1, 3], [2, 3], [0, 1]], int8=True) + + def test_ep2_empty_batch(self): + self._check_routing([]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_qwen36.py b/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_qwen36.py new file mode 100644 index 000000000000..b801189f2122 --- /dev/null +++ b/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_qwen36.py @@ -0,0 +1,235 @@ +"""Exercise main/draft MoE routing on two Ascend NPUs with Qwen3.6. + +Example (from the repository root): + ASCEND_RT_VISIBLE_DEVICES=0,1 \ + SGLANG_TEST_MODEL_PATH=/path/to/Qwen3.6-35B-A3B \ + python test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_qwen36.py -v + +SGLANG_TEST_TARGET_A2A and SGLANG_TEST_DRAFT_A2A default to none. The default +covers the NPU none dispatcher at EP=2, both without a draft and with MTP. +Draft none inherits the target backend, just like an omitted draft option. +Set SGLANG_TEST_TARGET_A2A=deepep to check inheritance of DeepEP instead. +Set SGLANG_TEST_EP_SIZE=1 to exercise the unsharded-expert TP path. +SGLANG_TEST_LOG_DIR preserves the server logs and response/metric artifacts. +Uses the launch/cleanup pattern in python/sglang/test/ascend. +""" + +import json +import os +import tempfile +import unittest +from pathlib import Path + +import requests +import torch +from transformers import AutoTokenizer + +from sglang.srt.utils import is_npu, kill_process_tree +from sglang.srt.utils.network import get_open_port +from sglang.test.ascend.test_ascend_utils import MODEL_WEIGHTS_DIR +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.test_utils import CustomTestCase, popen_launch_server + +register_npu_ci(est_time=600, suite="full-2-npu-a3", nightly=True) + + +@unittest.skipUnless(is_npu(), "Requires Ascend NPU") +class TestNpuSpeculativeMoeA2AQwen36(CustomTestCase): + def setUp(self): + self.assertGreaterEqual(torch.npu.device_count(), 2) + self.model = os.environ.get( + "SGLANG_TEST_MODEL_PATH", + os.path.join(MODEL_WEIGHTS_DIR, "Qwen/Qwen3.6-35B-A3B"), + ) + self.target_backend = os.environ.get("SGLANG_TEST_TARGET_A2A", "none") + self.draft_backend = os.environ.get("SGLANG_TEST_DRAFT_A2A", "none") + self.ep_size = int(os.environ.get("SGLANG_TEST_EP_SIZE", "2")) + log_dir = os.environ.get("SGLANG_TEST_LOG_DIR") + if log_dir is None: + temp_dir = tempfile.TemporaryDirectory(prefix="sglang-npu-spec-a2a-") + self.addCleanup(temp_dir.cleanup) + log_dir = temp_dir.name + self.log_dir = Path(log_dir) + self.log_dir.mkdir(parents=True, exist_ok=True) + self.session = requests.Session() + self.session.trust_env = False + self.addCleanup(self.session.close) + tokenizer = AutoTokenizer.from_pretrained(self.model, trust_remote_code=True) + self.answers = ["42", "56", "81", "15"] + self.prompts = [ + tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": f"Calculate {expression}. Reply with only the integer answer.", + } + ], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + for expression in ["19 + 23", "7 * 8", "9 * 9", "45 / 3"] + ] + + def _run_backend(self, label, draft_backend, *, speculative=True): + base_url = f"http://127.0.0.1:{get_open_port()}" + args = [ + "--device", + "npu", + "--attention-backend", + "ascend", + "--trust-remote-code", + "--dtype", + "bfloat16", + "--mamba-ssm-dtype", + "bfloat16", + "--tp-size", + "2", + "--ep-size", + str(self.ep_size), + "--moe-a2a-backend", + self.target_backend, + "--mem-fraction-static", + "0.75", + "--context-length", + "2048", + "--max-total-tokens", + "4096", + "--max-running-requests", + "4", + "--max-mamba-cache-size", + "8", + "--chunked-prefill-size", + "1024", + "--disable-radix-cache", + "--disable-cuda-graph", + "--random-seed", + "0", + ] + if speculative: + args += [ + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "1", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "2", + "--speculative-draft-model-quantization", + "unquant", + "--speculative-draft-attention-backend", + "ascend", + "--speculative-moe-runner-backend", + "auto", + ] + if draft_backend is not None: + args += ["--speculative-moe-a2a-backend", draft_backend] + if "ascend_fuseep" in (self.target_backend, draft_backend): + # DeepEP FuseEP mode 2 accepts INT8 weights only. Probe the BF16 + # checkpoint with mode 1, which exposes a BF16 kernel interface. + args += ["--fuseep-mode", "1"] + env = { + **os.environ, + # Qwen3.6 has 128 local experts at EP=2. DeepEP normal-mode + # metadata alone needs more than the usual 200 MB test setting. + "HCCL_BUFFSIZE": "512", + "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "0", + # The A3 low-latency window supports at most 512. This test has + # only four concurrent requests and two tokens per verify pass. + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "128", + "TOKENIZERS_PARALLELISM": "false", + } + log_path = self.log_dir / f"{label}.log" + with log_path.open("w", encoding="utf-8") as log_file: + process = None + try: + process = popen_launch_server( + self.model, + base_url, + timeout=900, + other_args=args, + env=env, + return_stdout_stderr=(log_file, log_file), + ) + info_response = self.session.get(f"{base_url}/server_info", timeout=30) + info_response.raise_for_status() + info = info_response.json() + self.assertEqual(info["tp_size"], 2) + self.assertEqual(info["ep_size"], self.ep_size) + self.assertEqual( + info["speculative_algorithm"], "EAGLE" if speculative else None + ) + self.assertEqual(info["moe_a2a_backend"], self.target_backend) + self.assertEqual( + info["speculative_moe_a2a_backend"], + draft_backend if speculative else None, + ) + if not speculative: + self.assertIsNone(info["speculative_draft_model_path"]) + outputs = [] + # Cover a single request and a batch with distinct lengths/content. + for prompts, answers in [ + (self.prompts[:1], self.answers[:1]), + (self.prompts, self.answers), + ]: + response = self.session.post( + f"{base_url}/generate", + json={ + "text": prompts, + "sampling_params": {"temperature": 0, "max_new_tokens": 64}, + }, + timeout=180, + ) + response.raise_for_status() + results = response.json() + self.assertEqual(len(results), len(prompts)) + outputs.extend(results) + # Save evidence before assertions, including failing output. + (self.log_dir / f"{label}.json").write_text( + json.dumps({"server_info": info, "outputs": outputs}, indent=2), + encoding="utf-8", + ) + for result, answer in zip(results, answers): + self.assertEqual(result["text"].strip(), answer, result) + meta = result["meta_info"] + self.assertGreater(meta["completion_tokens"], 0) + if speculative: + self.assertGreater(meta.get("spec_verify_ct", 0), 0, meta) + self.assertGreater( + meta.get("spec_num_proposed_drafts", 0), 0, meta + ) + self.assertGreater( + meta.get("spec_accepted_drafts", 0), 0, meta + ) + else: + self.assertEqual(meta.get("spec_verify_ct", 0), 0, meta) + self.assertEqual( + meta.get("spec_num_proposed_drafts", 0), 0, meta + ) + return outputs + except Exception: + log_file.flush() + print(log_path.read_text(encoding="utf-8", errors="replace")[-16000:]) + raise + finally: + if process is not None: + kill_process_tree(process.pid) + process.wait(timeout=30) + + def test_main_only(self): + self._run_backend("main_only", None, speculative=False) + + def test_draft_backend_override(self): + inherited = self._run_backend("inherited", None) + overridden = self._run_backend( + f"draft_{self.draft_backend}", self.draft_backend + ) + self.assertEqual( + [output["text"] for output in inherited], + [output["text"] for output in overridden], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 7c763f52bf82..03df7d641998 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -3168,12 +3168,15 @@ def test_speculative_draft_backend_rejected(self): validate_deepep_v2_speculative_draft(args) def test_inherited_speculative_draft_backend_rejected(self): - args = self._args( - moe_runner_backend="deep_gemm", - speculative_algorithm="EAGLE", - ) - with self.assertRaisesRegex(ValueError, "speculative draft backend"): - validate_deepep_v2_speculative_draft(args) + for draft_backend in (None, "none"): + with self.subTest(draft_backend=draft_backend): + args = self._args( + moe_runner_backend="deep_gemm", + speculative_algorithm="EAGLE", + speculative_moe_a2a_backend=draft_backend, + ) + with self.assertRaisesRegex(ValueError, "speculative draft backend"): + validate_deepep_v2_speculative_draft(args) def test_ngram_does_not_inherit_a_draft_backend(self): args = self._args( diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 83d5f207a839..d3e83be90f23 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -1108,6 +1108,27 @@ def test_initialize_materializes_group(self): self.assertTrue(is_tbo_enabled()) self.assertEqual(get_flags().moe.deepep_config, "") + def test_speculative_none_inherits_main_backend(self): + from sglang.srt.layers.moe.utils import ( + get_moe_a2a_backend, + get_speculative_moe_a2a_backend, + speculative_moe_a2a_backend_context, + ) + + for main_backend in ("none", "deepep", "ascend_fuseep"): + for draft_backend in (None, "none"): + with self.subTest(main=main_backend, draft=draft_backend): + self._init( + moe_a2a_backend=main_backend, + speculative_moe_a2a_backend=draft_backend, + ) + self.assertEqual( + get_speculative_moe_a2a_backend().value, main_backend + ) + with speculative_moe_a2a_backend_context(): + self.assertEqual(get_moe_a2a_backend().value, main_backend) + self.assertEqual(get_moe_a2a_backend().value, main_backend) + def test_speculative_swap_and_restore(self): from sglang.srt.layers.moe.utils import ( get_moe_a2a_backend, @@ -1120,11 +1141,11 @@ def test_speculative_swap_and_restore(self): moe_a2a_backend="deepep", moe_runner_backend="triton", speculative_moe_runner_backend="auto", - speculative_moe_a2a_backend="none", + speculative_moe_a2a_backend="ascend_fuseep", ) with speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): self.assertEqual(get_moe_runner_backend().name, "AUTO") - self.assertTrue(get_moe_a2a_backend().is_none()) + self.assertTrue(get_moe_a2a_backend().is_ascend_fuseep()) # MTP layers are unquantized: fp4 allgather is forced off self.assertTrue(get_flags().moe.disable_fp4_allgather) self.assertTrue(get_flags().moe.speculative_context)