diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index d8c843485..dbec58605 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -680,6 +680,11 @@ class TboForwardBatchPreparer: ): output_dict[key] = None continue + elif key == "rids" and len(old_value) != num_seqs: + output_dict[key] = old_value[ + start_seq_index : min(end_seq_index, len(old_value)) + ] + continue assert ( len(old_value) == num_seqs ), f"{key=} {old_value=} {num_seqs=} {batch=}" diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b0c8bb870..c9e24145c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1339,9 +1339,7 @@ class DeepseekV2MoE(nn.Module): return None def op_gate(self, state): - if is_non_idle_and_non_empty( - state.forward_batch.forward_mode, state.hidden_states_mlp_input - ): + if state.hidden_states_mlp_input.shape[0] > 0: # router_logits: (num_tokens, n_experts) state.router_logits = self.gate(state.hidden_states_mlp_input) else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b22e8f0bc..89a497bcc 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7858,11 +7858,6 @@ class ServerArgs: "When enabling two batch overlap, moe_a2a_backend cannot be 'none'." ) - if self.enable_two_batch_overlap and self.enforce_shared_experts_fusion: - raise ValueError( - "--enable-two-batch-overlap and --enforce-shared-experts-fusion cannot be used together." - ) - # Check communications compression if self.enable_quant_communications and self.tp_size == 1: raise ValueError( diff --git a/test/registered/ep/test_tbo_shared_experts_fusion.py b/test/registered/ep/test_tbo_shared_experts_fusion.py new file mode 100644 index 000000000..265b436bf --- /dev/null +++ b/test/registered/ep/test_tbo_shared_experts_fusion.py @@ -0,0 +1,73 @@ +import os +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=900, stage="extra-b", runner_config="deepep-8-gpu-h200") + +DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324" + + +class TestTBOWithSharedExpertsFusion(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V3_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp", + "8", + "--enable-dp-attention", + "--dp", + "8", + "--moe-dense-tp-size", + "1", + "--moe-a2a-backend", + "deepep", + "--enable-two-batch-overlap", + "--enforce-shared-experts-fusion", + "--disable-cuda-graph", + "--max-running-requests", + "512", + ], + env={ + **os.environ, + "SGLANG_TBO_DEBUG": "1", + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + self.assertGreater(metrics["score"], 0.60) + + +if __name__ == "__main__": + unittest.main()