diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index 4cfcf6e83..4dcb9219c 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import random from concurrent import futures from typing import TYPE_CHECKING, List @@ -18,6 +19,7 @@ if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler GRAMMAR_TIMEOUT = envs.SGLANG_GRAMMAR_TIMEOUT.get() +GRAMMAR_POLL_INTERVAL = envs.SGLANG_GRAMMAR_POLL_INTERVAL.get() logger = logging.getLogger(__name__) @@ -36,6 +38,11 @@ class GrammarManager: else: self.grammar_backend = None + self.grammar_sync_group = scheduler.dp_tp_cpu_group + self.grammar_sync_size = scheduler.dp_tp_group.world_size + self.grammar_sync_entry = scheduler.dp_tp_group.first_rank + self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank + def __len__(self): return len(self.grammar_queue) @@ -94,70 +101,89 @@ class GrammarManager: return add_to_grammar_queue + def finalize_grammar(self, req: Req): + self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy()) + if req.grammar is INVALID_GRAMMAR_OBJ: + error_msg = f"Invalid grammar request: {req.grammar_key=}" + req.set_finish_with_abort(error_msg) + def get_ready_grammar_requests(self) -> List[Req]: """Move requests whose grammar objects are ready from grammar_queue to waiting_queue.""" - num_ready_reqs = 0 - num_timeout_reqs = 0 + num_timeout_pt = 0 + sim_timeout_prob = ( + -1 + if self.is_grammar_sync_entry + else envs.SGLANG_GRAMMAR_SIMULATE_TIMEOUT.get() + ) # Entry rank never simulates timeout + timeout_ct = GRAMMAR_TIMEOUT / GRAMMAR_POLL_INTERVAL + for req in self.grammar_queue: try: if req.finished(): # It is aborted by AbortReq num_ready_reqs += 1 continue - req.grammar = req.grammar.result(timeout=0.03) - self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy()) - if req.grammar is INVALID_GRAMMAR_OBJ: - error_msg = f"Invalid grammar request: {req.grammar_key=}" - req.set_finish_with_abort(error_msg) + if sim_timeout_prob > 0 and random.random() < sim_timeout_prob: + # Simulate timeout for non-entry ranks in TP sync group for testing + logger.warning( + f"Simulating grammar timeout on {self.scheduler.tp_rank=}" + ) + raise futures._base.TimeoutError() + req.grammar = req.grammar.result(timeout=GRAMMAR_POLL_INTERVAL) + self.finalize_grammar(req) num_ready_reqs += 1 + num_timeout_pt = num_ready_reqs except futures._base.TimeoutError: req.grammar_wait_ct += 1 # NOTE(lianmin): this timeout is the waiting time of the above line. It is # not the waiting time from it enters the grammar queue. - if req.grammar_wait_ct > GRAMMAR_TIMEOUT / 0.03: - num_timeout_reqs = 1 + if sim_timeout_prob > 0 or req.grammar_wait_ct > timeout_ct: + num_timeout_pt = num_ready_reqs + 1 break - if self.server_args.enable_dp_attention: - tp_size = self.scheduler.attn_tp_size - tp_group = self.scheduler.attn_tp_cpu_group - else: - tp_size = self.scheduler.tp_size - tp_group = self.scheduler.tp_cpu_group - - if tp_size > 1: + if self.grammar_sync_size > 1: # Sync across TP ranks to make sure they have the same number of ready requests - tensor = torch.tensor([num_ready_reqs, num_timeout_reqs], dtype=torch.int32) - torch.distributed.all_reduce( - tensor, op=torch.distributed.ReduceOp.MAX, group=tp_group + tensor = torch.tensor( + [num_ready_reqs, -num_ready_reqs, -num_timeout_pt], dtype=torch.int32 ) - num_ready_reqs_max, num_timeout_reqs_max = tensor.tolist() - - for i in range(num_ready_reqs, num_ready_reqs_max): - req = self.grammar_queue[i] - if req.finished(): # It is aborted by AbortReq - continue - req.grammar = req.grammar.result() - self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy()) - if req.grammar is INVALID_GRAMMAR_OBJ: - error_msg = f"Invalid grammar request: {req.grammar_key=}" - req.set_finish_with_abort(error_msg) + torch.distributed.all_reduce( + tensor, op=torch.distributed.ReduceOp.MIN, group=self.grammar_sync_group + ) + num_ready_reqs_min = tensor[0].item() + num_ready_reqs_max = -tensor[1].item() + num_timeout_pt_max = -tensor[2].item() else: + num_ready_reqs_min = num_ready_reqs num_ready_reqs_max = num_ready_reqs - num_timeout_reqs_max = num_timeout_reqs + num_timeout_pt_max = num_timeout_pt - for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_max): + if envs.SGLANG_GRAMMAR_SIMULATE_TIMEOUT.get() > 0: + # NOTE: in simulation timeout mode, if only some TP ranks report a timeout, + # we still treat those requests as ready instead of canceling them. + for i in range(num_ready_reqs_min, num_ready_reqs_max): + req = self.grammar_queue[i] + if req.finished(): + continue + if isinstance(req.grammar, futures.Future): + req.grammar = req.grammar.result() + self.finalize_grammar(req) + + num_ready_reqs_min = num_ready_reqs_max + + # NOTE: in non-simulation mode, cancel any request that times out on a subset of + # TP ranks because timeouts can diverge across ranks. + for i in range(num_ready_reqs_min, num_timeout_pt_max): req = self.grammar_queue[i] - req.grammar.cancel() - self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ) - error_msg = f"Grammar preprocessing timed out for {req.grammar_key=}" - req.set_finish_with_abort(error_msg) + if req.finished(): + continue + if isinstance(req.grammar, futures.Future): + req.grammar.cancel() + req.grammar = INVALID_GRAMMAR_OBJ + self.finalize_grammar(req) - num_ready_reqs = num_ready_reqs_max + num_timeout_reqs_max - - ready_grammar_reqs = self.grammar_queue[:num_ready_reqs] - self.grammar_queue = self.grammar_queue[num_ready_reqs:] + ready_grammar_reqs = self.grammar_queue[:num_timeout_pt_max] + self.grammar_queue = self.grammar_queue[num_timeout_pt_max:] return ready_grammar_reqs diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8f93a8e5f..1e3cf8c61 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -171,6 +171,10 @@ class Envs: SGLANG_IS_IN_CI_AMD = EnvBool(False) SGLANG_TEST_MAX_RETRY = EnvInt(None) + # Constrained Decoding (Grammar) + SGLANG_GRAMMAR_SIMULATE_TIMEOUT = EnvFloat(-1) + SGLANG_GRAMMAR_POLL_INTERVAL = EnvFloat(0.03) + # Test & Debug SGLANG_DETECT_SLOW_RANK = EnvBool(False) SGLANG_TEST_STUCK_DETOKENIZER = EnvFloat(0) diff --git a/test/registered/distributed/test_dp_attention_large.py b/test/registered/distributed/test_dp_attention_large.py new file mode 100644 index 000000000..591237e11 --- /dev/null +++ b/test/registered/distributed/test_dp_attention_large.py @@ -0,0 +1,136 @@ +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.kits.ebnf_constrained_kit import TestEBNFConstrainedMixin +from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin +from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + DEFAULT_MODEL_NAME_FOR_TEST_MLA, + DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_amd_ci, + popen_launch_server, +) + +register_cuda_ci(est_time=350, suite="stage-c-test-large-4-gpu") + + +class TestDPAttentionDP2TP4( + CustomTestCase, + TestJSONConstrainedMixin, + TestEBNFConstrainedMixin, + TestRegexConstrainedMixin, +): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + with envs.SGLANG_GRAMMAR_SIMULATE_TIMEOUT.override(0.5): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp=4", + "--enable-dp-attention", + "--dp=2", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_mgsm_en(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mgsm_en", + num_examples=None, + num_threads=1024, + ) + + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["score"], 0.8) + + +class TestDPAttentionDP2TP2DeepseekV3MTP( + CustomTestCase, + TestJSONConstrainedMixin, + TestEBNFConstrainedMixin, + TestRegexConstrainedMixin, +): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--trust-remote-code", + "--disable-radix", + "--speculative-algorithm=EAGLE", + "--speculative-num-steps=2", + "--speculative-eagle-topk=4", + "--speculative-num-draft-tokens=4", + "--speculative-draft-model-path", + DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN, + "--tp-size=4", + "--enable-dp-attention", + "--dp-size=2", + ] + if not is_in_amd_ci(): + other_args += ["--mem-frac", "0.7"] + with envs.SGLANG_GRAMMAR_SIMULATE_TIMEOUT.override(0.5): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + requests.get(self.base_url + "/flush_cache") + + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(metrics) + + self.assertGreater(metrics["accuracy"], 0.60) + + server_info = requests.get(self.base_url + "/get_server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + print( + f"###test_gsm8k (deepseek-v3 mtp + dp):\n" + f"accuracy={metrics['accuracy']=:.3f}\n" + f"{avg_spec_accept_length=:.3f}\n" + ) + self.assertGreater(avg_spec_accept_length, 2.5) + + +if __name__ == "__main__": + unittest.main()