From 76b3c698d6fe620fc0abfc4aac1a5b2a41464cf0 Mon Sep 17 00:00:00 2001 From: Raayan Dhar <58057652+raayandhar@users.noreply.github.com> Date: Fri, 9 Jan 2026 08:38:32 -0800 Subject: [PATCH] feat: reduce constrained-decoding overhead in TP (#13947) Signed-off-by: Raayan Dhar raayan.dhar@gmail.com Signed-off-by: raayandhar --- python/sglang/srt/layers/sampler.py | 23 +- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/managers/scheduler.py | 198 +++++++++++++----- .../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, 229 insertions(+), 71 deletions(-) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 55bef5652..d0f2b17e5 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -40,10 +40,13 @@ class Sampler(nn.Module): def __init__(self): super().__init__() self.use_nan_detection = get_global_server_args().enable_nan_detection - self.tp_sync_group = get_tp_group().device_group - - if is_dp_attention_enabled(): - self.tp_sync_group = get_attention_tp_group().device_group + 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 def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo @@ -205,13 +208,21 @@ class Sampler(nn.Module): batch_next_token_ids, ] - if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars: + 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: # 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 81ce51f7c..ab4fffb44 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1311,7 +1311,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.grammar for req in reqs), + has_grammar=any(req.sampling_params.has_grammar_constraint 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), @@ -2039,7 +2039,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.token_ids_logprobs = None self.has_stream = any(req.stream for req in self.reqs) - self.has_grammar = any(req.grammar for req in self.reqs) + self.has_grammar = any( + req.sampling_params.has_grammar_constraint 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 77263a738..a37ddaaee 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1532,40 +1532,52 @@ class Scheduler( self._add_request_to_queue(req) return - # Init grammar cache for this request + # 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. add_to_grammar_queue = False - 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) + 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) else: - 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) + key = ("structural_tag", req.sampling_params.structural_tag) - 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 - add_to_grammar_queue = True + 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: - if value is INVALID_GRAMMAR_OBJ: # We hit a cached invalid grammar. + value, cache_hit = self.grammar_backend.get_cached_or_future_value( + key, req.require_reasoning + ) + req.grammar = value + req.grammar_key = key + + if cache_hit and value is INVALID_GRAMMAR_OBJ: 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: @@ -2351,21 +2363,47 @@ 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: - error_msg = f"Invalid grammar request: {req.grammar_key=}" - req.set_finish_with_abort(error_msg) + invalid_indices.append(num_ready_reqs) num_ready_reqs += 1 except futures._base.TimeoutError: @@ -2376,44 +2414,90 @@ class Scheduler( num_timeout_reqs = 1 break - 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 - 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 + # 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 ) - 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) + # 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=}" + ) else: - num_ready_reqs_max = num_ready_reqs - num_timeout_reqs_max = num_timeout_reqs + num_ready_reqs_rank_0 = num_ready_reqs + num_timeout_reqs_rank_0 = num_timeout_reqs - for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_max): + # 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=}" + ) + + # 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 + ): req = self.grammar_queue[i] - req.grammar.cancel() - self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ) + # 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) error_msg = f"Grammar preprocessing timed out for {req.grammar_key=}" req.set_finish_with_abort(error_msg) - num_ready_reqs = num_ready_reqs_max + num_timeout_reqs_max + num_ready_reqs = num_ready_reqs_rank_0 + num_timeout_reqs_rank_0 + # 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 1be58b15c..30fc63220 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2300,10 +2300,19 @@ 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. - sampling_info.update_regex_vocab_mask() - sampling_info.apply_logits_bias(logits_output.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, + ) def sample( self, diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index a8f17c754..1bcc0e62b 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -201,7 +201,11 @@ class SamplingBatchInfo: return # Find a grammar from the list - first_grammar = next(grammar for grammar in self.grammars if grammar) + 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 # TODO(lianmin): Maybe we can reuse the existing mask? self.vocab_mask = first_grammar.allocate_vocab_mask( @@ -232,7 +236,7 @@ class SamplingBatchInfo: else: self.acc_linear_penalties = None - def apply_logits_bias(self, logits: torch.Tensor): + def apply_logits_bias(self, logits: torch.Tensor, apply_vocab_mask: bool = True): if self.acc_linear_penalties is not None: # Used in the overlap mode logits.add_(self.acc_linear_penalties) @@ -241,7 +245,7 @@ class SamplingBatchInfo: # Used in the non-overlap mode self.penalizer_orchestrator.apply(logits) - if self.vocab_mask is not None: + if apply_vocab_mask and 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 a7f4c664f..2da2f69e2 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -102,6 +102,19 @@ 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 d83f0d3e2..600b533e3 100644 --- a/test/registered/constrained_decoding/test_constrained_decoding.py +++ b/test/registered/constrained_decoding/test_constrained_decoding.py @@ -5,6 +5,8 @@ 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 @@ -21,6 +23,7 @@ from sglang.test.test_utils import ( class ServerWithGrammar(CustomTestCase): backend = "xgrammar" disable_overlap = False + tp_size = 1 @classmethod def setUpClass(cls): @@ -33,6 +36,9 @@ 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"] @@ -70,5 +76,34 @@ 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()