fix grammar timeout sync across tp ranks. (#16898)

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)
+                if sim_timeout_prob > 0 and random.random() < sim_timeout_prob:
-                self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
+                    # Simulate timeout for non-entry ranks in TP sync group for testing
-                if req.grammar is INVALID_GRAMMAR_OBJ:
+                    logger.warning(
-                    error_msg = f"Invalid grammar request: {req.grammar_key=}"
+                        f"Simulating grammar timeout on {self.scheduler.tp_rank=}"
-                    req.set_finish_with_abort(error_msg)
+                    )
+                    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:
+                if sim_timeout_prob > 0 or req.grammar_wait_ct > timeout_ct:
-                    num_timeout_reqs = 1
+                    num_timeout_pt = num_ready_reqs + 1
                 break
 
-        if self.server_args.enable_dp_attention:
+        if self.grammar_sync_size > 1:
-            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:
             # 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)
+            tensor = torch.tensor(
-            torch.distributed.all_reduce(
+                [num_ready_reqs, -num_ready_reqs, -num_timeout_pt], dtype=torch.int32
-                tensor, op=torch.distributed.ReduceOp.MAX, group=tp_group
             )
-            num_ready_reqs_max, num_timeout_reqs_max = tensor.tolist()
+            torch.distributed.all_reduce(
-
+                tensor, op=torch.distributed.ReduceOp.MIN, group=self.grammar_sync_group
-            for i in range(num_ready_reqs, num_ready_reqs_max):
+            )
-                req = self.grammar_queue[i]
+            num_ready_reqs_min = tensor[0].item()
-                if req.finished():  # It is aborted by AbortReq
+            num_ready_reqs_max = -tensor[1].item()
-                    continue
+            num_timeout_pt_max = -tensor[2].item()
-                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)
         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()
+            if req.finished():
-            self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ)
+                continue
-            error_msg = f"Grammar preprocessing timed out for {req.grammar_key=}"
+            if isinstance(req.grammar, futures.Future):
-            req.set_finish_with_abort(error_msg)
+                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_timeout_pt_max]
-
+        self.grammar_queue = self.grammar_queue[num_timeout_pt_max:]
-        ready_grammar_reqs = self.grammar_queue[:num_ready_reqs]
-        self.grammar_queue = self.grammar_queue[num_ready_reqs:]
 
         return ready_grammar_reqs

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)

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()