From 578b119bc66746d096a14917245af93c503674a1 Mon Sep 17 00:00:00 2001 From: Siyuan Chen <41201609+SYChen123@users.noreply.github.com> Date: Sat, 31 Jan 2026 03:57:42 +0800 Subject: [PATCH] [BugFix] Fix server crashes when req.grammar and ngram spec are enabled (#17585) --- python/sglang/srt/speculative/ngram_info.py | 3 ++ python/sglang/srt/speculative/ngram_worker.py | 32 ++++++++++++++++++- python/sglang/srt/speculative/spec_utils.py | 13 +++++--- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 67194aef6..210ee505c 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -7,6 +7,7 @@ from typing import Optional, Tuple import torch import triton +from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject from sglang.srt.server_args import get_global_server_args logger = logging.getLogger(__name__) @@ -57,6 +58,7 @@ class NgramVerifyInput(SpecInput): retrive_next_token: torch.Tensor, retrive_next_sibling: torch.Tensor, draft_token_num: int, + grammar: BaseGrammarObject = None, ): super().__init__(SpecInputType.NGRAM_VERIFY) self.draft_token = draft_token @@ -67,6 +69,7 @@ class NgramVerifyInput(SpecInput): self.retrive_next_sibling = retrive_next_sibling self.draft_token_num = draft_token_num self.device = self.custom_mask.device + self.grammar = grammar def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]: return self.draft_token_num, self.draft_token_num diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index fdeec130d..0a830fd95 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -14,6 +14,7 @@ from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.cpp_ngram.ngram_cache import NgramCache from sglang.srt.speculative.ngram_info import NgramVerifyInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_utils import generate_token_bitmask logger = logging.getLogger(__name__) @@ -213,10 +214,18 @@ class NGRAMWorker: def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult: self._prepare_for_speculative_decoding(batch) model_worker_batch = batch.get_model_worker_batch() + spec_info = model_worker_batch.spec_info num_accepted_tokens = 0 accept_lens = None if model_worker_batch.forward_mode.is_target_verify(): + if batch.has_grammar: + retrieve_next_token_cpu = spec_info.retrive_next_token.cpu() + retrieve_next_sibling_cpu = spec_info.retrive_next_sibling.cpu() + draft_tokens_cpu = spec_info.draft_token.view( + spec_info.retrive_next_token.shape + ).cpu() + batch_result = self.target_worker.forward_batch_generation( model_worker_batch, is_verify=True ) @@ -224,9 +233,30 @@ class NGRAMWorker: batch_result.logits_output, batch_result.can_run_cuda_graph, ) + verify_input: NgramVerifyInput = model_worker_batch.spec_info + vocab_mask = None + if batch.has_grammar: + # Generate the logit mask for structured output. + # Overlap the CPU operations for bitmask generation with the forward pass. + vocab_mask = generate_token_bitmask( + batch.reqs, + verify_input, + retrieve_next_token_cpu, + retrieve_next_sibling_cpu, + draft_tokens_cpu, + batch.sampling_info.vocab_size, + ) + + if vocab_mask is not None: + assert verify_input.grammar is not None + vocab_mask = vocab_mask.to(verify_input.retrive_next_token.device) + # NOTE (sk): otherwise, this vocab mask will be the one from the previous extend stage + # and will be applied to produce wrong results + batch.sampling_info.vocab_mask = None + logits_output, next_token_ids, num_accepted_tokens = verify_input.verify( - batch, logits_output, self.page_size + batch, logits_output, self.page_size, vocab_mask ) # Store accept_lens for per-request metrics accept_lens = verify_input.accept_length diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 2e8057e8f..39e2806f5 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -573,6 +573,7 @@ def traverse_tree( draft_tokens: torch.Tensor, grammar: BaseGrammarObject, allocate_token_bitmask: torch.Tensor, + vocab_size: Optional[int] = None, ): """ Traverse the tree constructed by the draft model to generate the logits mask. @@ -594,10 +595,13 @@ def traverse_tree( else: parent_bitmask = allocate_token_bitmask[parent_pos] curr_token_id = draft_tokens[curr] - # 32 boolean bitmask values are packed into 32-bit integers - accepted = ( - parent_bitmask[curr_token_id // 32] & (1 << (curr_token_id % 32)) - ) != 0 + if vocab_size and curr_token_id >= vocab_size: + accepted = False + else: + # 32 boolean bitmask values are packed into 32-bit integers + accepted = ( + parent_bitmask[curr_token_id // 32] & (1 << (curr_token_id % 32)) + ) != 0 if accepted: if curr != 0: @@ -670,6 +674,7 @@ def generate_token_bitmask( allocate_token_bitmask[ i * num_draft_tokens : (i + 1) * num_draft_tokens ], + vocab_size=vocab_size, ) tree_traverse_time = time.perf_counter() - s if tree_traverse_time > TREE_TRAVERSE_TIME_THRESHOLD: