[BugFix] Fix server crashes when req.grammar and ngram spec are enabled (#17585)

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

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

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: