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

Revert "feat: reduce constrained-decoding overhead in TP" (#16845)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-01-10 11:39:38 +08:00
committed by GitHub
parent fbc128a32e
commit e6d40bff81
7 changed files with 73 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/sglang/srt/layers/sampler.py
View File

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

6
python/sglang/srt/managers/schedule_batch.py
View File

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

202
python/sglang/srt/managers/scheduler.py
View File

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

15
python/sglang/srt/model_executor/model_runner.py
View File

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

10
python/sglang/srt/sampling/sampling_batch_info.py
View File

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

13
python/sglang/srt/sampling/sampling_params.py
View File

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