From e6d40bff81644416b2fa2fd193136ee446ccde53 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 10 Jan 2026 11:39:38 +0800 Subject: [PATCH] Revert "feat: reduce constrained-decoding overhead in TP" (#16845) --- python/sglang/srt/layers/sampler.py | 23 +- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/managers/scheduler.py | 202 +++++------------- .../sglang/srt/model_executor/model_runner.py | 15 +- .../srt/sampling/sampling_batch_info.py | 10 +- python/sglang/srt/sampling/sampling_params.py | 13 -- .../test_constrained_decoding.py | 35 --- 7 files changed, 73 insertions(+), 231 deletions(-) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index d0f2b17e5..55bef5652 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -40,13 +40,10 @@ class Sampler(nn.Module): def __init__(self): super().__init__() self.use_nan_detection = get_global_server_args().enable_nan_detection - tp_group = ( - get_attention_tp_group() if is_dp_attention_enabled() else get_tp_group() - ) - self.tp_rank = tp_group.rank_in_group - self.tp_size = tp_group.world_size - self.tp_root_rank = tp_group.ranks[0] - self.tp_sync_group = tp_group.device_group + self.tp_sync_group = get_tp_group().device_group + + if is_dp_attention_enabled(): + self.tp_sync_group = get_attention_tp_group().device_group def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo @@ -208,21 +205,13 @@ class Sampler(nn.Module): batch_next_token_ids, ] - if sampling_info.grammars: - if self.tp_size > 1: - # Grammar-aware sampling only runs on TP rank 0 - # We broadcast its choice to keep ranks in sync. - dist.broadcast( - batch_next_token_ids, - src=self.tp_root_rank, - group=self.tp_sync_group, - ) - elif SYNC_TOKEN_IDS_ACROSS_TP: + if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: # For performance reasons, SGLang does not sync the final token IDs across TP ranks by default. # This saves one all-reduce, but the correctness of this approach depends on the determinism of several operators: # the last all-reduce, the last lm_head matmul, and all sampling kernels. # These kernels are deterministic in most cases, but there are some rare instances where they are not deterministic. # In such cases, enable this env variable to prevent hanging due to TP ranks becoming desynchronized. + # When using xgrammar, this becomes more likely so we also do the sync when grammar is used. torch.distributed.all_reduce( batch_next_token_ids, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4cfa135e1..5f168595c 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1313,7 +1313,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): enable_overlap=enable_overlap, return_logprob=return_logprob, has_stream=any(req.stream for req in reqs), - has_grammar=any(req.sampling_params.has_grammar_constraint for req in reqs), + has_grammar=any(req.grammar for req in reqs), device=req_to_token_pool.device, spec_algorithm=spec_algorithm, return_hidden_states=any(req.return_hidden_states for req in reqs), @@ -2041,9 +2041,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.token_ids_logprobs = None self.has_stream = any(req.stream for req in self.reqs) - self.has_grammar = any( - req.sampling_params.has_grammar_constraint for req in self.reqs - ) + self.has_grammar = any(req.grammar for req in self.reqs) self.sampling_info.filter_batch(keep_indices, keep_indices_device) # NOTE: spec_info filtered before batch filtering only happens in: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ce466dc94..49088634a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1531,52 +1531,40 @@ class Scheduler( self._add_request_to_queue(req) return - # When TP>1 only rank 0 compiles grammars to avoid duplicating work across all ranks. - # All ranks add requests to the grammar_queue for synchronization. Each rank still - # maintains its own cache populated with ready grammars. Rank 0 broadcasts readiness. + # Init grammar cache for this request add_to_grammar_queue = False - if req.sampling_params.has_grammar_constraint: - if req.sampling_params.json_schema is not None: - key = ("json", req.sampling_params.json_schema) - elif req.sampling_params.regex is not None: - key = ("regex", req.sampling_params.regex) - elif req.sampling_params.ebnf is not None: - key = ("ebnf", req.sampling_params.ebnf) + if ( + req.sampling_params.json_schema is not None + or req.sampling_params.regex is not None + or req.sampling_params.ebnf is not None + or req.sampling_params.structural_tag is not None + ): + if self.grammar_backend is None: + error_msg = "Grammar-based generation (json_schema, regex, ebnf, structural_tag) is not supported when the server is launched with --grammar-backend none" + req.set_finish_with_abort(error_msg) else: - key = ("structural_tag", req.sampling_params.structural_tag) + if req.sampling_params.json_schema is not None: + key = ("json", req.sampling_params.json_schema) + elif req.sampling_params.regex is not None: + key = ("regex", req.sampling_params.regex) + elif req.sampling_params.ebnf is not None: + key = ("ebnf", req.sampling_params.ebnf) + elif req.sampling_params.structural_tag: + key = ("structural_tag", req.sampling_params.structural_tag) - if self.tp_rank == 0: - # Only rank 0 compiles grammar and manages the cache - if self.server_args.grammar_backend == "none": - error_msg = "Grammar-based generation (json_schema, regex, ebnf, structural_tag) is not supported when the server is launched with --grammar-backend none" - req.set_finish_with_abort(error_msg) - if self.tp_size > 1: - add_to_grammar_queue = True - else: - value, cache_hit = self.grammar_backend.get_cached_or_future_value( - key, req.require_reasoning - ) - req.grammar = value + value, cache_hit = self.grammar_backend.get_cached_or_future_value( + key, req.require_reasoning + ) + req.grammar = value + + if not cache_hit: req.grammar_key = key - - if cache_hit and value is INVALID_GRAMMAR_OBJ: + add_to_grammar_queue = True + else: + if value is INVALID_GRAMMAR_OBJ: # We hit a cached invalid grammar. error_msg = f"Invalid grammar request with cache hit: {key=}" req.set_finish_with_abort(error_msg) - # When TP > 1, always add to grammar_queue for synchronization across ranks - if self.tp_size > 1: - add_to_grammar_queue = True - elif not cache_hit: - # TP == 1: only add to queue if we didn't get a cache hit. - add_to_grammar_queue = True - else: - # Non-rank-0 workers add to grammar queue for synchronization, but no compilation. - if self.server_args.grammar_backend == "none": - error_msg = "Grammar-based generation (json_schema, regex, ebnf, structural_tag) is not supported when the server is launched with --grammar-backend none" - req.set_finish_with_abort(error_msg) - req.grammar_key = key - add_to_grammar_queue = True - if add_to_grammar_queue: self.grammar_queue.append(req) else: @@ -2362,47 +2350,21 @@ class Scheduler( self.send_to_tokenizer.send_output(HealthCheckOutput()) def move_ready_grammar_requests(self): - """ - Move requests whose grammar objects are ready from grammar_queue to waiting_queue. + """Move requests whose grammar objects are ready from grammar_queue to waiting_queue.""" - When TP>1, rank 0 counts the compiled or timed out grammars and broadcasts this to - all other ranks. All ranks process the same number of requests and non-entry workers - fetch the compiled grammars from their local cache. - """ num_ready_reqs = 0 num_timeout_reqs = 0 - invalid_indices = [] - - if self.server_args.enable_dp_attention: - tp_size = self.attn_tp_size - tp_group = self.attn_tp_cpu_group - else: - tp_size = self.tp_size - tp_group = self.tp_cpu_group - - # Count how many requests have ready grammars (non-blocking check). - # For TP>1, only rank 0 actually checks Future completion; others just count. for req in self.grammar_queue: try: if req.finished(): # It is aborted by AbortReq num_ready_reqs += 1 continue - if tp_size > 1: - # TP>1: Non-rank-0 workers have req.grammar=None (they don't compile). - # They still count requests as "ready" to stay in sync with rank 0. - # Rank 0 with cache hits will have a compiled grammar (not a Future). - # Only rank 0 with cache misses has Futures that need .result() below. - if req.grammar is None or not isinstance( - req.grammar, futures.Future - ): - 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: - invalid_indices.append(num_ready_reqs) + error_msg = f"Invalid grammar request: {req.grammar_key=}" + req.set_finish_with_abort(error_msg) num_ready_reqs += 1 except futures._base.TimeoutError: @@ -2413,90 +2375,44 @@ class Scheduler( num_timeout_reqs = 1 break - if tp_size > 1: - # Rank 0 broadcasts its ready/timeout/invalid counts - num_invalid = len(invalid_indices) - tensor = torch.tensor( - [num_ready_reqs, num_timeout_reqs, num_invalid], dtype=torch.int32 - ) - - # Broadcast from local rank 0 within the TP group, using group_src (local rank) - # rather than src (global rank) to be consistent with other self.tp_rank == 0 checks - torch.distributed.broadcast(tensor, group_src=0, group=tp_group) - num_ready_reqs_rank_0, num_timeout_reqs_rank_0, num_invalid = ( - tensor.tolist() - ) - - # In order to have the consistent abort request state, we must broadcast the invalid indices - if num_invalid > 0: - if self.tp_rank == 0: - # On rank 0, we create our tensor with our invalid indices data (to send) - invalid_tensor = torch.tensor(invalid_indices, dtype=torch.int32) - else: - # On non-entry ranks, we create our empty tensors (to receive) - invalid_tensor = torch.zeros(num_invalid, dtype=torch.int32) - torch.distributed.broadcast(invalid_tensor, group_src=0, group=tp_group) - - # All ranks abort invalid requests - for idx in invalid_tensor.tolist(): - req = self.grammar_queue[idx] - if not req.finished(): - req.set_finish_with_abort( - f"Invalid grammar request: {req.grammar_key=}" - ) + if self.server_args.enable_dp_attention: + tp_size = self.attn_tp_size + tp_group = self.attn_tp_cpu_group else: - num_ready_reqs_rank_0 = num_ready_reqs - num_timeout_reqs_rank_0 = num_timeout_reqs + tp_size = self.tp_size + tp_group = self.tp_cpu_group - # Non TP>1: Handle invalid grammars directly - for idx in invalid_indices: - req = self.grammar_queue[idx] - req.set_finish_with_abort( - f"Invalid grammar request: {req.grammar_key=}" - ) + 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) + torch.distributed.all_reduce( + tensor, op=torch.distributed.ReduceOp.MAX, group=tp_group + ) + num_ready_reqs_max, num_timeout_reqs_max = tensor.tolist() - # Handle timed-out requests: all ranks must abort to maintain consistency. - for i in range( - num_ready_reqs_rank_0, num_ready_reqs_rank_0 + num_timeout_reqs_rank_0 - ): + 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) + else: + num_ready_reqs_max = num_ready_reqs + num_timeout_reqs_max = num_timeout_reqs + + for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_max): req = self.grammar_queue[i] - # Only rank 0 has futures to cancel and cache to update - if self.tp_rank == 0 and isinstance(req.grammar, futures.Future): - req.grammar.cancel() - self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ) + 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) - num_ready_reqs = num_ready_reqs_rank_0 + num_timeout_reqs_rank_0 + num_ready_reqs = num_ready_reqs_max + num_timeout_reqs_max - # Move synchronized requests from grammar_queue to waiting_queue. for req in self.grammar_queue[:num_ready_reqs]: - # Skip cache lookup for already-finished requests (invalid/timed-out/aborted) - if ( - not req.finished() - and req.grammar is None - and self.grammar_backend is not None - and req.grammar_key is not None - ): - # TP>1: Non-rank-0 workers fetch grammar from rank 0's cache. - grammar_obj, cache_hit = ( - self.grammar_backend.get_cached_or_future_value( - req.grammar_key, req.require_reasoning - ) - ) - - # If we got a Future, wait for compilation to complete. - if isinstance(grammar_obj, futures.Future): - grammar_obj = grammar_obj.result() - self.grammar_backend.set_cache(req.grammar_key, grammar_obj.copy()) - - if grammar_obj is INVALID_GRAMMAR_OBJ: - req.set_finish_with_abort( - f"Invalid grammar request: {req.grammar_key=}" - ) - else: - req.grammar = grammar_obj - self._add_request_to_queue(req) self.grammar_queue = self.grammar_queue[num_ready_reqs:] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 30fc63220..1be58b15c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2300,19 +2300,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): ): # NOTE: In overlap mode, the function update_regex_vocab_mask (in sample) # was executed after we processed last batch's results. - has_grammar = sampling_info.grammars is not None - # Avoid compiling grammar masks on every TP rank; only the first rank applies them. - apply_vocab_mask = has_grammar and self.sampler.tp_rank == 0 + # Calculate logits bias and apply it to next_token_logits. - if apply_vocab_mask: - sampling_info.update_regex_vocab_mask() - else: - sampling_info.vocab_mask = None - sampling_info.apply_mask_func = None - sampling_info.apply_logits_bias( - logits_output.next_token_logits, - apply_vocab_mask=apply_vocab_mask, - ) + sampling_info.update_regex_vocab_mask() + sampling_info.apply_logits_bias(logits_output.next_token_logits) def sample( self, diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 1bcc0e62b..a8f17c754 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -201,11 +201,7 @@ class SamplingBatchInfo: return # Find a grammar from the list - first_grammar = next((grammar for grammar in self.grammars if grammar), None) - if first_grammar is None: - self.vocab_mask = None - self.apply_mask_func = None - return + first_grammar = next(grammar for grammar in self.grammars if grammar) # TODO(lianmin): Maybe we can reuse the existing mask? self.vocab_mask = first_grammar.allocate_vocab_mask( @@ -236,7 +232,7 @@ class SamplingBatchInfo: else: self.acc_linear_penalties = None - def apply_logits_bias(self, logits: torch.Tensor, apply_vocab_mask: bool = True): + def apply_logits_bias(self, logits: torch.Tensor): if self.acc_linear_penalties is not None: # Used in the overlap mode logits.add_(self.acc_linear_penalties) @@ -245,7 +241,7 @@ class SamplingBatchInfo: # Used in the non-overlap mode self.penalizer_orchestrator.apply(logits) - if apply_vocab_mask and self.vocab_mask is not None: + if self.vocab_mask is not None: self.apply_mask_func(logits=logits, vocab_mask=self.vocab_mask) if self.logit_bias is not None: diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 2da2f69e2..a7f4c664f 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -102,19 +102,6 @@ class SamplingParams: if self.top_k == -1: self.top_k = TOP_K_ALL # whole vocabulary - @property - def has_grammar_constraint(self) -> bool: - """ - Helper property to check if any grammar constraint - exists for these SamplingParams. - """ - return ( - self.json_schema is not None - or self.regex is not None - or self.ebnf is not None - or self.structural_tag is not None - ) - def verify(self, vocab_size): if self.temperature < 0.0: raise ValueError( diff --git a/test/registered/constrained_decoding/test_constrained_decoding.py b/test/registered/constrained_decoding/test_constrained_decoding.py index 600b533e3..d83f0d3e2 100644 --- a/test/registered/constrained_decoding/test_constrained_decoding.py +++ b/test/registered/constrained_decoding/test_constrained_decoding.py @@ -5,8 +5,6 @@ register_amd_ci(est_time=179, suite="stage-b-test-small-1-gpu-amd") import unittest -import torch - from sglang.srt.utils import kill_process_tree from sglang.test.kits.ebnf_constrained_kit import TestEBNFConstrainedMinxin from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin @@ -23,7 +21,6 @@ from sglang.test.test_utils import ( class ServerWithGrammar(CustomTestCase): backend = "xgrammar" disable_overlap = False - tp_size = 1 @classmethod def setUpClass(cls): @@ -36,9 +33,6 @@ class ServerWithGrammar(CustomTestCase): cls.backend, ] - if cls.tp_size > 1: - launch_args += ["--tp-size", str(cls.tp_size)] - if cls.disable_overlap: launch_args += ["--disable-overlap-schedule"] @@ -76,34 +70,5 @@ class TestLLGuidanceBackend( backend = "llguidance" -# Test with TP=2 -@unittest.skipIf(torch.cuda.device_count() < 2, "Requires at least 2 GPUs") -class TestXGrammarBackendTP2( - ServerWithGrammar, - TestJSONConstrainedMixin, - TestEBNFConstrainedMinxin, - TestRegexConstrainedMixin, -): - backend = "xgrammar" - tp_size = 2 - - -@unittest.skipIf(torch.cuda.device_count() < 2, "Requires at least 2 GPUs") -class TestOutlinesBackendTP2(ServerWithGrammar, TestJSONConstrainedMixin): - backend = "outlines" - tp_size = 2 - - -@unittest.skipIf(torch.cuda.device_count() < 2, "Requires at least 2 GPUs") -class TestLLGuidanceBackendTP2( - ServerWithGrammar, - TestJSONConstrainedMixin, - TestEBNFConstrainedMinxin, - TestRegexConstrainedMixin, -): - backend = "llguidance" - tp_size = 2 - - if __name__ == "__main__": unittest.main()