Use native batched llguidance mask generation (#32412)

Co-authored-by: Alec S <10566873+alecsolder@users.noreply.github.com>

python/sglang/srt/constrained/base_grammar_backend.py

@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""The baseclass of a backend for grammar-guided constrained decoding."""
+"""The base class of a backend for grammar-guided constrained decoding."""
 
 import logging
 import time
@@ -22,10 +22,13 @@ from typing import Dict, List, NamedTuple, Optional, Tuple
 import torch
 
 from sglang.srt.parser.reasoning_parser import ReasoningParser
+from sglang.srt.runtime_context import get_resources
 from sglang.srt.server_args import ServerArgs
 
 logger = logging.getLogger(__name__)
 
+GRAMMAR_BACKEND_REGISTRY = {}
+
 
 @dataclass
 class GrammarStats:
@@ -39,6 +42,13 @@ class GrammarStats:
     num_timeout: int = 0
 
 
+class GrammarRow(NamedTuple):
+    """Grammar and destination row for a batched vocab-mask fill."""
+
+    row: int
+    grammar: "BaseGrammarObject"
+
+
 class BaseGrammarObject:
 
     def __init__(self):
@@ -69,6 +79,19 @@ class BaseGrammarObject:
     def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
         raise NotImplementedError()
 
+    @staticmethod
+    def fill_vocab_mask_batched(
+        entries: List[GrammarRow], vocab_mask: torch.Tensor
+    ) -> None:
+        """Fill listed rows, leaving unlisted rows untouched."""
+        for entry in entries:
+            entry.grammar.fill_vocab_mask(vocab_mask, entry.row)
+
+    @staticmethod
+    def reset_vocab_mask(vocab_mask: torch.Tensor) -> None:
+        """Restore a reusable mask to the backend's unconstrained state."""
+        raise NotImplementedError()
+
     @staticmethod
     def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
         raise NotImplementedError()
@@ -148,6 +171,16 @@ class BaseGrammarBackend:
         self.executor = ThreadPoolExecutor()
         self.cache: Dict[Tuple[str, str], BaseGrammarObject] = {}
 
+    def initialize_vocab_mask_buffer(
+        self,
+        name: str,
+        vocab_size: int,
+        max_rows: int,
+        device,
+    ) -> Optional[torch.Tensor]:
+        """Initialize a reusable mask buffer when supported by the backend."""
+        return None
+
     def _not_supported(self, key_type: str, key_string: str) -> BaseGrammarObject:
         logger.warning(f"Skip unsupported {key_type=}, {key_string=}")
         return InvalidGrammarObject()
@@ -226,7 +259,49 @@ class BaseGrammarBackend:
         self.cache.clear()
 
 
-GRAMMAR_BACKEND_REGISTRY = {}
+def register_vocab_mask_buffer(
+    name: str, vocab_mask: torch.Tensor, max_rows: int
+) -> torch.Tensor:
+    """Register a fixed-capacity mask buffer, preserving an equivalent one."""
+    if max_rows <= 0:
+        raise ValueError(f"Grammar mask max_rows must be positive, got {max_rows}")
+    if vocab_mask.ndim == 0 or vocab_mask.shape[0] != max_rows:
+        raise ValueError(
+            f"Grammar mask buffer {name!r} must have {max_rows} rows, "
+            f"got shape {tuple(vocab_mask.shape)}"
+        )
+
+    buffers = get_resources().buffers
+    existing = buffers.get(name)
+    if existing is not None:
+        if (
+            existing.shape != vocab_mask.shape
+            or existing.dtype != vocab_mask.dtype
+            or existing.device != vocab_mask.device
+        ):
+            raise RuntimeError(
+                f"Grammar mask buffer {name!r} was already initialized as "
+                f"{tuple(existing.shape)}, {existing.dtype}, {existing.device}; "
+                f"new buffer is {tuple(vocab_mask.shape)}, {vocab_mask.dtype}, "
+                f"{vocab_mask.device}"
+            )
+        return existing
+
+    buffers[name] = vocab_mask
+    return vocab_mask
+
+
+def get_vocab_mask_buffer(name: str, rows: int) -> Optional[torch.Tensor]:
+    """Return the active rows of a registered mask buffer, if available."""
+    vocab_mask = get_resources().buffers.get(name)
+    if vocab_mask is None:
+        return None
+    if rows > vocab_mask.shape[0]:
+        raise ValueError(
+            f"Grammar batch needs {rows} mask rows, exceeding initialized "
+            f"capacity {vocab_mask.shape[0]} for {name!r}"
+        )
+    return vocab_mask[:rows]
 
 
 def register_grammar_backend(name, init_func):

python/sglang/srt/constrained/llguidance_backend.py

@@ -15,26 +15,76 @@
 
 import json
 import logging
-import os
-from typing import Iterable, List, Optional, Tuple, Union
+from functools import cache
+from typing import Iterable, List, NamedTuple, Optional, Tuple, Union
 
 import torch
-from llguidance import LLMatcher, LLTokenizer, StructTag, grammar_from
+from llguidance import LLExecutor, LLMatcher, LLTokenizer, StructTag, grammar_from
 from llguidance.hf import from_tokenizer
 from llguidance.torch import (
     allocate_token_bitmask,
     apply_token_bitmask_inplace,
     fill_next_token_bitmask,
+    fill_next_token_bitmask_par,
+    fill_next_token_bitmask_par_with_draft_tokens,
 )
 
 from sglang.srt.constrained.base_grammar_backend import (
     BaseGrammarBackend,
     BaseGrammarObject,
+    GrammarRow,
     InvalidGrammarObject,
+    register_vocab_mask_buffer,
 )
 from sglang.srt.constrained.utils import is_legacy_structural_tag
+from sglang.srt.utils import get_int_env_var
 
 logger = logging.getLogger(__name__)
+_LLGUIDANCE_LOG_LEVEL = get_int_env_var("LLGUIDANCE_LOG_LEVEL", 1)
+
+
+class GrammarDraftRow(NamedTuple):
+    """Grammar, destination block, and tokens for a draft-chain mask fill."""
+
+    base_row: int
+    grammar: "GuidanceGrammar"
+    draft_tokens: List[int]
+
+
+@cache
+def _get_or_init_mask_executor() -> LLExecutor:
+    return LLExecutor()
+
+
+def fill_token_bitmask_with_draft_tokens(
+    entries: List[GrammarDraftRow],
+    vocab_mask: torch.Tensor,
+) -> None:
+    """Fill speculative draft-chain masks with llguidance's native kernel.
+
+    Each matcher is advanced over its legal draft prefix and rolled back by the
+    number of consumed parser tokens before returning. Rows after the first
+    illegal token retain the caller's all-allow value. Correctness requires
+    matcher rollback to restore parser-token state exactly. Matcher errors retain
+    the native executor's behavior; this wrapper adds no per-matcher Python probe.
+    """
+    if not entries:
+        return
+    matchers = [(e.grammar.ll_matcher, e.base_row, e.draft_tokens) for e in entries]
+    fill_next_token_bitmask_par_with_draft_tokens(
+        _get_or_init_mask_executor(), matchers, vocab_mask
+    )
+
+
+def fill_token_bitmask_batched(
+    entries: List[GrammarRow],
+    vocab_mask: torch.Tensor,
+) -> None:
+    """Fill regular-decode mask rows with llguidance's native kernel."""
+    if not entries:
+        return
+    matchers = [(e.grammar.ll_matcher, e.row) for e in entries]
+    fill_next_token_bitmask_par(_get_or_init_mask_executor(), matchers, vocab_mask)
 
 
 def _normalize_eos_token_ids(
@@ -47,15 +97,26 @@ def _normalize_eos_token_ids(
 
 class GuidanceGrammar(BaseGrammarObject):
 
-    def __init__(self, llguidance_tokenizer: LLTokenizer, serialized_grammar: str):
+    def __init__(
+        self,
+        llguidance_tokenizer: LLTokenizer,
+        serialized_grammar: str,
+        *,
+        ll_matcher: Optional[LLMatcher] = None,
+    ):
         super().__init__()
         self.llguidance_tokenizer = llguidance_tokenizer
         self.serialized_grammar = serialized_grammar
 
-        self.ll_matcher = LLMatcher(
-            self.llguidance_tokenizer,
-            self.serialized_grammar,
-            log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")),
+        # A request copy reuses the cached template's compiled matcher.
+        self.ll_matcher = (
+            ll_matcher
+            if ll_matcher is not None
+            else LLMatcher(
+                self.llguidance_tokenizer,
+                self.serialized_grammar,
+                log_level=_LLGUIDANCE_LOG_LEVEL,
+            )
         )
         self._check_err()
 
@@ -87,6 +148,24 @@ class GuidanceGrammar(BaseGrammarObject):
         fill_next_token_bitmask(self.ll_matcher, vocab_mask, idx)
         self._check_err()
 
+    @staticmethod
+    def fill_vocab_mask_batched(
+        entries: List[GrammarRow], vocab_mask: torch.Tensor
+    ) -> None:
+        """Use the native fill when every entry is a plain llguidance grammar."""
+        if all(isinstance(entry.grammar, GuidanceGrammar) for entry in entries):
+            fill_token_bitmask_batched(entries, vocab_mask)
+            return
+        BaseGrammarObject.fill_vocab_mask_batched(entries, vocab_mask)
+
+    @staticmethod
+    def reset_vocab_mask(vocab_mask: torch.Tensor) -> None:
+        if vocab_mask.dtype != torch.int32:
+            raise TypeError(
+                f"llguidance requires a packed int32 mask, got {vocab_mask.dtype}"
+            )
+        vocab_mask.fill_(-1)
+
     def allocate_vocab_mask(
         self, vocab_size: int, batch_size: int, device
     ) -> torch.Tensor:
@@ -101,9 +180,12 @@ class GuidanceGrammar(BaseGrammarObject):
         apply_token_bitmask_inplace(logits, vocab_mask)
 
     def copy(self):
+        # Cache templates are pristine, so cloning their matcher creates a fresh
+        # request grammar without recompiling the serialized grammar.
         return GuidanceGrammar(
             llguidance_tokenizer=self.llguidance_tokenizer,
             serialized_grammar=self.serialized_grammar,
+            ll_matcher=self.ll_matcher.deep_copy(),
         )
 
     def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
@@ -146,6 +228,21 @@ class GuidanceBackend(BaseGrammarBackend):
             n_vocab,
             eos_token=_normalize_eos_token_ids(eos_token_ids),
         )
+        # Initialize the shared executor here so the first batched mask fill
+        # does not pay its one-time setup cost on the request path.
+        _get_or_init_mask_executor()
+
+    def initialize_vocab_mask_buffer(
+        self,
+        name: str,
+        vocab_size: int,
+        max_rows: int,
+        device,
+    ) -> torch.Tensor:
+        vocab_mask = allocate_token_bitmask(
+            max_rows, self.llguidance_tokenizer.vocab_size
+        )
+        return register_vocab_mask_buffer(name, vocab_mask, max_rows)
 
     def _from_serialized(self, serialized_grammar) -> BaseGrammarObject:
         try:

python/sglang/srt/sampling/sampling_batch_info.py

@@ -10,6 +10,7 @@ import sglang.srt.sampling.penaltylib as penaltylib
 from sglang.srt.constrained.base_grammar_backend import (
     BaseGrammarObject,
     GrammarMask,
+    GrammarRow,
 )
 from sglang.srt.runtime_context import get_server_args
 from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
@@ -244,17 +245,20 @@ class SamplingBatchInfo:
         # 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?
         vocab_mask = first_grammar.allocate_vocab_mask(
             vocab_size=self.vocab_size,
             batch_size=len(self.temperatures),
             device=self.device,
         )
 
-        # 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(vocab_mask, i)
+        # Rows omitted here (finished / terminated / non-grammar requests) retain
+        # the freshly allocated buffer's unconstrained value.
+        entries = [
+            GrammarRow(row=row, grammar=grammar)
+            for row, grammar in enumerate(self.grammars)
+            if grammar and not grammar.finished and not grammar.is_terminated()
+        ]
+        first_grammar.fill_vocab_mask_batched(entries, vocab_mask)
 
         # Move the mask to the device if needed
         vocab_mask = first_grammar.move_vocab_mask(vocab_mask, self.device)