[Spec] Hold the grammar bitmask in one GrammarMask type across all decode paths (#32409)

python/sglang/srt/constrained/base_grammar_backend.py

@@ -17,7 +17,7 @@ import logging
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, NamedTuple, Optional, Tuple
 
 import torch
 
@@ -117,6 +117,19 @@ class BaseGrammarObject:
         raise NotImplementedError()
 
 
+class GrammarMask(NamedTuple):
+    """A filled vocab_mask plus the backend that applies it.
+
+    The grammar is any one of the batch's -- a handle, not per-request state.
+    """
+
+    grammar: BaseGrammarObject
+    vocab_mask: torch.Tensor
+
+    def apply(self, logits: torch.Tensor) -> None:
+        self.grammar.apply_vocab_mask(logits=logits, vocab_mask=self.vocab_mask)
+
+
 class InvalidGrammarObject(BaseGrammarObject):
     """Represents a grammar that failed to compile, carrying the original error message."""
 

python/sglang/srt/model_executor/model_runner.py

@@ -1591,10 +1591,10 @@ class ModelRunner:
 
         # Release the vocab_mask GPU tensor immediately after it has been applied
         # to the logits. In overlap scheduling, the sampling_info (and its
-        # vocab_mask) can be kept alive by the delay_sample_func closure and
+        # grammar_mask) can be kept alive by the delay_sample_func closure and
         # batch_record_buf until the next iteration, causing a steady VRAM leak
         # when structured output (grammar) is used.
-        sampling_info.vocab_mask = None
+        sampling_info.grammar_mask = None
 
     def sample(
         self,

python/sglang/srt/sampling/sampling_batch_info.py

@@ -2,11 +2,15 @@ from __future__ import annotations
 
 import dataclasses
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import torch
 
 import sglang.srt.sampling.penaltylib as penaltylib
+from sglang.srt.constrained.base_grammar_backend import (
+    BaseGrammarObject,
+    GrammarMask,
+)
 from sglang.srt.runtime_context import get_server_args
 from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
 from sglang.srt.sampling.penaltylib.repetition_penalty import apply_scaling_penalties
@@ -44,11 +48,10 @@ class SamplingBatchInfo:
 
     # Masking tensors for grammar-guided structured outputs
     vocab_size: int
-    grammars: Optional[List] = None
+    grammars: Optional[List[Optional[BaseGrammarObject]]] = None
     rids_int: Optional[torch.Tensor] = None
     bootstrap_room_ids_int: Optional[torch.Tensor] = None
-    vocab_mask: Optional[torch.Tensor] = None
+    grammar_mask: Optional[GrammarMask] = None
-    apply_mask_func: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None
 
     # Penalizer
     penalizer_orchestrator: Optional[penaltylib.BatchedPenalizerOrchestrator] = None
@@ -235,30 +238,27 @@ class SamplingBatchInfo:
 
     def update_regex_vocab_mask(self):
         if not self.grammars:
-            self.vocab_mask = None
+            self.grammar_mask = None
-            self.apply_mask_func = None
             return
 
         # Find a grammar from the list
         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(
+        vocab_mask = first_grammar.allocate_vocab_mask(
             vocab_size=self.vocab_size,
             batch_size=len(self.temperatures),
             device=self.device,
         )
-        self.apply_mask_func = (
-            first_grammar.apply_vocab_mask
-        )  # force to use static method
 
         # Apply the mask
         for i, grammar in enumerate(self.grammars):
             if grammar and not grammar.finished and not grammar.is_terminated():
-                grammar.fill_vocab_mask(self.vocab_mask, i)
+                grammar.fill_vocab_mask(vocab_mask, i)
 
         # Move the mask to the device if needed
-        self.vocab_mask = first_grammar.move_vocab_mask(self.vocab_mask, self.device)
+        vocab_mask = first_grammar.move_vocab_mask(vocab_mask, self.device)
+        self.grammar_mask = GrammarMask(first_grammar, vocab_mask)
 
     def update_penalties(self):
         if self.penalizer_orchestrator.is_required:
@@ -290,8 +290,8 @@ class SamplingBatchInfo:
             # Used in the non-overlap mode
             self.penalizer_orchestrator.apply(logits)
 
-        if self.vocab_mask is not None:
+        if self.grammar_mask is not None:
-            self.apply_mask_func(logits=logits, vocab_mask=self.vocab_mask)
+            self.grammar_mask.apply(logits)
 
         if self.logit_bias is not None:
             logits.add_(self.logit_bias)

python/sglang/srt/speculative/dflash_info.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Optional
 import torch
 
 from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton
-from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
@@ -45,9 +44,6 @@ class DFlashVerifyInput(SpecInput):
 
     ragged_verify_layout: Optional[RaggedVerifyLayout] = None
 
-    # Stamped by generate_token_bitmask during verify, read back to apply the mask.
-    grammar: Optional[BaseGrammarObject] = None
-
     def __post_init__(self):
         super().__init__(spec_input_type=SpecInputType.DFLASH_VERIFY)
         if self.num_tokens_per_req == -1:

python/sglang/srt/speculative/dflash_info_v2.py

@@ -6,7 +6,6 @@ from typing import List, Optional
 
 import torch
 
-from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.mem_cache.allocation import alloc_for_spec_decode
@@ -57,9 +56,6 @@ class DFlashDraftInputV2(SpecInput):
 
     verify_token_budget: Optional[int] = None
 
-    # Stamped by generate_token_bitmask during verify, read back to apply the mask.
-    grammar: Optional[BaseGrammarObject] = None
-
     def __post_init__(self):
         super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT)
         # Spec v2 draft state itself does not change token accounting.

python/sglang/srt/speculative/dflash_utils.py

@@ -145,7 +145,7 @@ def apply_dflash_verify_logits_adjustments(
 
     acc_linear_penalties = getattr(sampling_info, "acc_linear_penalties", None)
     penalizer = getattr(sampling_info, "penalizer_orchestrator", None)
-    vocab_mask = getattr(sampling_info, "vocab_mask", None)
+    grammar_mask = getattr(sampling_info, "grammar_mask", None)
     logit_bias = getattr(sampling_info, "logit_bias", None)
 
     logits_3d: Optional[torch.Tensor] = None
@@ -161,7 +161,7 @@ def apply_dflash_verify_logits_adjustments(
     # broadcast over the verify block without materializing a repeated buffer.
     if (
         penalizer is not None and penalizer.is_required and acc_linear_penalties is None
-    ) or vocab_mask is not None:
+    ) or grammar_mask is not None:
         linear_penalty = torch.zeros(
             (bs, next_token_logits.shape[1]),
             dtype=torch.float32,

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -1688,19 +1688,14 @@ class DFlashWorkerV2(BaseSpecWorker):
         logits_output = target_out.logits_output
         can_run_cuda_graph = target_out.can_run_cuda_graph
 
-        vocab_mask = None
+        grammar_mask = None
         if batch.has_grammar:
-            # Grammar barrier: advance the previous batch's FSM over its committed
+            grammar_mask = build_grammar_vocab_mask(
-            # tokens before building this batch's bitmask. Runs after the target
-            # launch, so the advance and the traversal both overlap the forward.
-            if grammar_barrier is not None:
-                grammar_barrier()
-            vocab_mask = build_grammar_vocab_mask(
                 reqs=batch.reqs,
-                verify_input=verify_input,
                 tree=grammar_tree,
                 sampling_info=batch.sampling_info,
                 device=logits_output.next_token_logits.device,
+                barrier=grammar_barrier,
             )
 
         if sampling_info is not None:
@@ -1711,10 +1706,8 @@ class DFlashWorkerV2(BaseSpecWorker):
             )
 
         # Constrain every chain position before accept picks from it.
-        if vocab_mask is not None:
+        if grammar_mask is not None:
-            verify_input.grammar.apply_vocab_mask(
+            grammar_mask.apply(logits_output.next_token_logits)
-                logits=logits_output.next_token_logits, vocab_mask=vocab_mask
-            )
 
         candidates = draft_tokens
         new_seq_lens = None

python/sglang/srt/speculative/dspark_components/dspark_verify.py

@@ -54,7 +54,7 @@ def verify_logits_adjustments_are_noop(sampling_info) -> bool:
     penalizer = getattr(sampling_info, "penalizer_orchestrator", None)
     if penalizer is not None and penalizer.is_required:
         return False
-    if getattr(sampling_info, "vocab_mask", None) is not None:
+    if getattr(sampling_info, "grammar_mask", None) is not None:
         return False
     if getattr(sampling_info, "logit_bias", None) is not None:
         return False

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -612,23 +612,17 @@ class DSparkWorkerV2(BaseSpecWorker):
         can_run_cuda_graph = target_verify.can_run_cuda_graph
 
         if batch.has_grammar:
-            # Both the FSM advance over the previous batch's committed tokens and
-            # the traversal below are host work, so they overlap the launch above.
-            if grammar_barrier is not None:
-                grammar_barrier()
             # run_compact scatters its rows back to (bs * chain_len), so the mask
             # lines up with the logits on both verify paths.
-            vocab_mask = build_grammar_vocab_mask(
+            grammar_mask = build_grammar_vocab_mask(
                 reqs=batch.reqs,
-                verify_input=draft_input,
                 tree=grammar_tree,
                 sampling_info=sampling_info,
                 device=logits_output.next_token_logits.device,
+                barrier=grammar_barrier,
             )
-            if vocab_mask is not None:
+            if grammar_mask is not None:
-                draft_input.grammar.apply_vocab_mask(
+                grammar_mask.apply(logits_output.next_token_logits)
-                    logits=logits_output.next_token_logits, vocab_mask=vocab_mask
-                )
 
         epilogue = self._verify_executor.verify_epilogue
         folded_accept = fold_eligible and run_compact and can_run_cuda_graph

python/sglang/srt/speculative/eagle_info.py

@@ -5,7 +5,6 @@ from typing import List, Optional
 import torch
 
 from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton
-from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
 from sglang.srt.runtime_context import get_server_args
 from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
@@ -28,7 +27,6 @@ class EagleVerifyInput(SpecInput):
     capture_hidden_mode: CaptureHiddenMode
     seq_lens_sum: int
     seq_lens_cpu: torch.Tensor
-    grammar: BaseGrammarObject = None
     # Stacked per-step draft proposal distribution q, shape (bs, num_steps,
     # vocab); only set under rejection sampling. Consumed by the verify kernel.
     draft_probs: torch.Tensor = None

python/sglang/srt/speculative/eagle_utils.py

@@ -28,6 +28,7 @@ from sglang.srt.utils import (
 from sglang.srt.utils.async_probe import maybe_detect_oob
 
 if TYPE_CHECKING:
+    from sglang.srt.constrained.base_grammar_backend import GrammarMask
     from sglang.srt.layers.logits_processor import LogitsProcessorOutput
     from sglang.srt.managers.schedule_batch import ScheduleBatch
     from sglang.srt.managers.tp_worker import TpModelWorker
@@ -641,7 +642,7 @@ def eagle_sample(
     verify_input: EagleVerifyInput,
     batch: ScheduleBatch,
     logits_output: LogitsProcessorOutput,
-    vocab_mask: torch.Tensor = None,
+    grammar_mask: Optional[GrammarMask] = None,
 ):
     """
     Verify and find accepted tokens based on logits output and batch
@@ -702,11 +703,8 @@ def eagle_sample(
         )
 
     # Apply grammar mask if provided
-    if vocab_mask is not None:
+    if grammar_mask is not None:
-        assert verify_input.grammar is not None
+        grammar_mask.apply(next_token_logits)
-        verify_input.grammar.apply_vocab_mask(
-            logits=next_token_logits, vocab_mask=vocab_mask
-        )
 
     candidates = verify_input.draft_token.reshape(bs, verify_input.draft_token_num)
     predict_shape = list(next_token_logits.shape)[:-1]

python/sglang/srt/speculative/eagle_worker_common.py

@@ -564,20 +564,14 @@ def run_eagle_verify(
     logits_output = forward_batch_output.logits_output
 
     # Generate vocab mask for constrained decoding
-    vocab_mask = None
+    grammar_mask = None
     if batch.has_grammar:
-        # Grammar barrier: advance the previous batch's grammar FSM over its
+        grammar_mask = build_grammar_vocab_mask(
-        # committed tokens before building this batch's bitmask. Runs after the
-        # target forward launch, so the FSM advance and the traversal below both
-        # overlap the target verify forward. No-op if there is nothing pending.
-        if grammar_barrier is not None:
-            grammar_barrier()
-        vocab_mask = build_grammar_vocab_mask(
             reqs=batch.reqs,
-            verify_input=verify_input,
             tree=grammar_tree,
             sampling_info=batch.sampling_info,
             device=verify_input.retrieve_next_token.device,
+            barrier=grammar_barrier,
         )
 
     # Sample
@@ -587,7 +581,7 @@ def run_eagle_verify(
         predict,
         accept_lens,
         accept_index,
-    ) = eagle_sample(verify_input, batch, logits_output, vocab_mask)
+    ) = eagle_sample(verify_input, batch, logits_output, grammar_mask)
     new_seq_lens = batch.seq_lens + accept_lens
     clear_unaccepted_c128 = getattr(
         token_to_kv_pool_allocator.get_kvcache(),

python/sglang/srt/speculative/ngram_info.py

@@ -5,7 +5,6 @@ from typing import List, Optional
 import torch
 
 from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton
-from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
 
 
@@ -19,7 +18,6 @@ class NgramVerifyInput(SpecInput):
         retrieve_next_token: torch.Tensor = None,
         retrieve_next_sibling: torch.Tensor = None,
         draft_token_num: int = None,
-        grammar: BaseGrammarObject = None,
         future_indices: Optional[torch.Tensor] = None,
         new_seq_lens: Optional[torch.Tensor] = None,
         accept_tokens: Optional[torch.Tensor] = None,
@@ -35,7 +33,6 @@ class NgramVerifyInput(SpecInput):
         self.draft_token_num = draft_token_num
         self.num_tokens_per_req = draft_token_num
         self.num_tokens_for_logprob_per_req = draft_token_num
-        self.grammar = grammar
 
         # Inputs for V2 overlap worker
         self.future_indices = future_indices

python/sglang/srt/speculative/ngram_worker.py

@@ -423,7 +423,7 @@ class NGRAMWorker(BaseSpecWorker):
             )
 
             verify_input: NgramVerifyInput = batch.spec_info
-            vocab_mask = None
+            grammar_mask = None
             if batch.has_grammar:
                 # From the host tree rather than the device output: no readback to
                 # wait on, and deriving here keeps it under the verify forward.
@@ -431,9 +431,8 @@ class NGRAMWorker(BaseSpecWorker):
                 retrieve_next_token_cpu, retrieve_next_sibling_cpu = _derive_tree_links(
                     mask, bs, self.draft_token_num
                 )
-                vocab_mask = build_grammar_vocab_mask(
+                grammar_mask = build_grammar_vocab_mask(
                     reqs=batch.reqs,
-                    verify_input=verify_input,
                     tree=GrammarTree.from_host(
                         retrieve_next_token_cpu,
                         retrieve_next_sibling_cpu,
@@ -441,6 +440,8 @@ class NGRAMWorker(BaseSpecWorker):
                     ),
                     sampling_info=batch.sampling_info,
                     device=verify_input.retrieve_next_token.device,
+                    # Host corpus lookup, so NGRAM stays synchronous: nothing pending.
+                    barrier=None,
                 )
 
             # Sample
@@ -454,7 +455,7 @@ class NGRAMWorker(BaseSpecWorker):
                 predict,
                 accept_lens,
                 accept_index,
-            ) = eagle_sample(verify_input, batch, logits_output, vocab_mask)
+            ) = eagle_sample(verify_input, batch, logits_output, grammar_mask)
             new_seq_lens = batch.seq_lens + accept_lens
             commit_mamba_states_after_verify(
                 self.target_worker,

python/sglang/srt/speculative/spec_utils.py

@@ -5,7 +5,7 @@ import logging
 import os
 import time
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, List, Literal, Optional, Tuple
 
 import torch
 from huggingface_hub import snapshot_download
@@ -32,6 +32,7 @@ from sglang.kernels.ops.speculative.eagle import (
     fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func,
 )
 from sglang.srt.configs.hybrid_arch import mambaish_config
+from sglang.srt.constrained.base_grammar_backend import GrammarMask
 from sglang.srt.distributed.parallel_state import (
     GroupCoordinator,
     patch_tensor_parallel_group,
@@ -72,8 +73,6 @@ if TYPE_CHECKING:
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
     from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
     from sglang.srt.server_args import ServerArgs
-    from sglang.srt.speculative.eagle_info import EagleVerifyInput
-    from sglang.srt.speculative.spec_info import SpecInput
 
 
 if _is_cuda:
@@ -502,12 +501,11 @@ def traverse_tree(
 
 def generate_token_bitmask(
     reqs: List[Req],
-    verify_input: EagleVerifyInput,
     retrieve_next_token_cpu: torch.Tensor,
     retrieve_next_sibling_cpu: torch.Tensor,
     draft_tokens_cpu: torch.Tensor,
     vocab_size: int,
-):
+) -> Tuple[Optional[torch.Tensor], Optional[BaseGrammarObject]]:
     """
     Generate the logit mask for structured output.
     Draft model's token can be either valid or invalid with respect to the grammar.
@@ -548,8 +546,7 @@ def generate_token_bitmask(
                     f"grammar: {req.grammar}"
                 )
 
-    verify_input.grammar = grammar
+    return allocate_token_bitmask, grammar
-    return allocate_token_bitmask
 
 
 class GrammarTree:
@@ -613,32 +610,33 @@ class GrammarTree:
 def build_grammar_vocab_mask(
     *,
     reqs: List[Req],
-    verify_input: SpecInput,
     tree: GrammarTree,
     sampling_info: SamplingBatchInfo,
     device,
-) -> Optional[torch.Tensor]:
+    barrier: Optional[Callable[[], None]],
+) -> Optional[GrammarMask]:
     """Build the constrained-decoding bitmask over a verify tree and stage it on device.
 
-    Call it after the target verify launch: resolving the tree and traversing it are
+    Call it after the target verify launch -- every step here is host work, so it all
-    both host work, so both overlap that forward.
+    overlaps that forward. ``barrier`` advances the previous batch's FSM over its
+    committed tokens, which the traversal then reads, so it has to run first.
     """
-    vocab_mask = generate_token_bitmask(
+    if barrier is not None:
+        barrier()
+    vocab_mask, grammar = generate_token_bitmask(
         reqs,
-        verify_input,
         *tree.resolve(),
         sampling_info.vocab_size,
     )
     if vocab_mask is None:
         return None
 
-    assert verify_input.grammar is not None
     # non_blocking is safe: the bitmask is pinned (see xgrammar_backend), and stream
     # order keeps the copy ahead of the sampler's apply_vocab_mask.
     vocab_mask = vocab_mask.to(device, non_blocking=True)
     # Otherwise the extend stage's leftover mask is applied instead.
-    sampling_info.vocab_mask = None
+    sampling_info.grammar_mask = None
-    return vocab_mask
+    return GrammarMask(grammar, vocab_mask)
 
 
 def load_token_map(token_map_path: str) -> List[int]:

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch
 
 import torch
 
+from sglang.srt.constrained.base_grammar_backend import GrammarMask
 from sglang.srt.sampling.sampling_batch_info import (
     SamplingBatchInfo,
     merge_bias_tensor,
@@ -164,13 +165,13 @@ class TestApplyLogitsBias(CustomTestCase):
         self.assertAlmostEqual(logits[0, 0].item(), 0.0, places=5)
 
     def test_applies_vocab_mask(self):
-        """Test that vocab_mask triggers the apply_mask_func callback."""
+        """Test that a grammar_mask gets applied to the logits."""
         info = _make_info(batch_size=1)
-        info.vocab_mask = torch.ones(1, VOCAB_SIZE)
+        grammar = MagicMock()
-        info.apply_mask_func = MagicMock()
+        info.grammar_mask = GrammarMask(grammar, torch.ones(1, VOCAB_SIZE))
         logits = torch.zeros(1, VOCAB_SIZE)
         info.apply_logits_bias(logits)
-        info.apply_mask_func.assert_called_once()
+        grammar.apply_vocab_mask.assert_called_once()
 
     def test_applies_penalizer_orchestrator(self):
         """Test that a required orchestrator's apply() is called on logits."""
@@ -185,7 +186,7 @@ class TestApplyLogitsBias(CustomTestCase):
         info = _make_info(batch_size=1)
         info.acc_additive_penalties = None
         info.logit_bias = None
-        info.vocab_mask = None
+        info.grammar_mask = None
         logits = torch.zeros(1, VOCAB_SIZE)
         original = logits.clone()
         info.apply_logits_bias(logits)
@@ -220,19 +221,18 @@ class TestUpdatePenalties(CustomTestCase):
 class TestUpdateRegexVocabMask(CustomTestCase):
 
     def test_no_grammars_clears_mask(self):
-        """Test that None grammars clears both vocab_mask and apply_mask_func."""
+        """Test that None grammars clears the grammar_mask."""
         info = _make_info(batch_size=1)
         info.grammars = None
         info.update_regex_vocab_mask()
-        self.assertIsNone(info.vocab_mask)
+        self.assertIsNone(info.grammar_mask)
-        self.assertIsNone(info.apply_mask_func)
 
     def test_empty_grammars_clears_mask(self):
-        """Test that empty grammars list clears vocab_mask."""
+        """Test that empty grammars list clears the grammar_mask."""
         info = _make_info(batch_size=1)
         info.grammars = []
         info.update_regex_vocab_mask()
-        self.assertIsNone(info.vocab_mask)
+        self.assertIsNone(info.grammar_mask)
 
     def test_with_grammars_allocates_and_fills(self):
         """Test that an active grammar gets allocate, fill, and move called."""
@@ -247,6 +247,7 @@ class TestUpdateRegexVocabMask(CustomTestCase):
         grammar.allocate_vocab_mask.assert_called_once()
         grammar.fill_vocab_mask.assert_called_once()
         grammar.move_vocab_mask.assert_called_once()
+        self.assertIs(info.grammar_mask.grammar, grammar)
 
     def test_mixed_grammars_only_active_fills(self):
         """Test that finished, terminated, and None grammars are skipped."""