chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-01-14 10:26:31 +08:00
committed by GitHub
parent a4825ed588
commit e2c8a50b38
3 changed files with 207 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
python/sglang/srt/constrained/grammar_manager.py
View File

@@ -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

4
python/sglang/srt/environ.py
View File

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

136
test/registered/distributed/test_dp_attention_large.py Normal file
View File

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