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 functools import cache
-from typing import Iterable, List, Optional, Tuple, Union
+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(
+        # A request copy reuses the cached template's compiled matcher.
-            self.llguidance_tokenizer,
+        self.ll_matcher = (
-            self.serialized_grammar,
+            ll_matcher
-            log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")),
+            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
+        # Rows omitted here (finished / terminated / non-grammar requests) retain
-        for i, grammar in enumerate(self.grammars):
+        # the freshly allocated buffer's unconstrained value.
-            if grammar and not grammar.finished and not grammar.is_terminated():
+        entries = [
-                grammar.fill_vocab_mask(vocab_mask, i)
+            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)

test/registered/unit/constrained/test_llguidance_batched_mask.py

@@ -0,0 +1,112 @@
+"""Bitwise parity tests for llguidance's regular-decode batched mask fill."""
+
+import unittest
+from unittest.mock import MagicMock
+
+import torch
+from llguidance import LLTokenizer, grammar_from
+
+from sglang.srt.constrained.base_grammar_backend import GrammarRow
+from sglang.srt.constrained.llguidance_backend import GuidanceBackend, GuidanceGrammar
+from sglang.srt.runtime_context import get_resources
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+_REGEX = r"[0-9]{1,8}"
+
+
+class TestLLGuidanceBatchedMask(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.template = GuidanceGrammar(
+            llguidance_tokenizer=LLTokenizer("byte"),
+            serialized_grammar=grammar_from("regex", _REGEX),
+        )
+
+    def _fresh(self, n):
+        return [self.template.copy() for _ in range(n)]
+
+    def _allocate(self, grammars):
+        return grammars[0].allocate_vocab_mask(
+            self.template.llguidance_tokenizer.vocab_size, len(grammars), "cpu"
+        )
+
+    def _serial(self, grammars):
+        mask = self._allocate(grammars)
+        for row, grammar in enumerate(grammars):
+            if not grammar.finished and not grammar.is_terminated():
+                grammar.fill_vocab_mask(mask, row)
+        return mask
+
+    def _batched(self, grammars):
+        mask = self._allocate(grammars)
+        entries = [
+            GrammarRow(row=row, grammar=grammar)
+            for row, grammar in enumerate(grammars)
+            if not grammar.finished and not grammar.is_terminated()
+        ]
+        grammars[0].fill_vocab_mask_batched(entries, mask)
+        return mask
+
+    def test_batched_matches_serial(self):
+        for batch_size in (1, 4, 10):
+            with self.subTest(batch_size=batch_size):
+                serial = self._serial(self._fresh(batch_size))
+                batched = self._batched(self._fresh(batch_size))
+                self.assertTrue(torch.equal(serial, batched))
+                self.assertTrue((batched[0] != -1).any())
+
+    def test_finished_row_stays_all_allow(self):
+        serial_grammars = self._fresh(3)
+        batched_grammars = self._fresh(3)
+        serial_grammars[1].finished = True
+        batched_grammars[1].finished = True
+
+        serial = self._serial(serial_grammars)
+        batched = self._batched(batched_grammars)
+
+        self.assertTrue(torch.equal(serial, batched))
+        self.assertTrue((batched[1] == -1).all())
+
+    def test_unsupported_entry_uses_serial_fill(self):
+        mask = self._allocate(self._fresh(1))
+        fallback = MagicMock()
+        fallback.fill_vocab_mask.side_effect = lambda vocab_mask, row: vocab_mask[
+            row
+        ].zero_()
+        entries = [GrammarRow(row=0, grammar=fallback)]
+
+        self.template.fill_vocab_mask_batched(entries, mask)
+
+        fallback.fill_vocab_mask.assert_called_once_with(mask, 0)
+        self.assertTrue((mask == 0).all())
+
+    def test_backend_initializes_fixed_mask_buffer(self):
+        name = "test_llguidance_vocab_mask"
+        get_resources().buffers.pop(name, None)
+        backend = object.__new__(GuidanceBackend)
+        backend.llguidance_tokenizer = self.template.llguidance_tokenizer
+
+        try:
+            mask = backend.initialize_vocab_mask_buffer(
+                name=name,
+                vocab_size=self.template.llguidance_tokenizer.vocab_size,
+                max_rows=4,
+                device="cpu",
+            )
+            same_mask = backend.initialize_vocab_mask_buffer(
+                name=name,
+                vocab_size=self.template.llguidance_tokenizer.vocab_size,
+                max_rows=4,
+                device="cpu",
+            )
+
+            self.assertEqual(mask.shape[0], 4)
+            self.assertEqual(mask.data_ptr(), same_mask.data_ptr())
+        finally:
+            get_resources().buffers.pop(name, None)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -43,6 +43,11 @@ def _make_info(batch_size=2, **overrides):
     return SamplingBatchInfo(**defaults)
 
 
+def _serial_batched_fill(entries, vocab_mask):
+    for entry in entries:
+        entry.grammar.fill_vocab_mask(vocab_mask, entry.row)
+
+
 class TestMergeBiasTensor(CustomTestCase):
 
     def test_both_none_returns_none(self):
@@ -239,12 +244,14 @@ class TestUpdateRegexVocabMask(CustomTestCase):
         grammar = MagicMock()
         grammar.finished = False
         grammar.is_terminated.return_value = False
+        grammar.fill_vocab_mask_batched.side_effect = _serial_batched_fill
         grammar.allocate_vocab_mask.return_value = torch.zeros(1, VOCAB_SIZE)
         grammar.move_vocab_mask.return_value = torch.zeros(1, VOCAB_SIZE)
         info = _make_info(batch_size=1)
         info.grammars = [grammar]
         info.update_regex_vocab_mask()
         grammar.allocate_vocab_mask.assert_called_once()
+        grammar.fill_vocab_mask_batched.assert_called_once()
         grammar.fill_vocab_mask.assert_called_once()
         grammar.move_vocab_mask.assert_called_once()
         self.assertIs(info.grammar_mask.grammar, grammar)
@@ -254,6 +261,7 @@ class TestUpdateRegexVocabMask(CustomTestCase):
         active = MagicMock()
         active.finished = False
         active.is_terminated.return_value = False
+        active.fill_vocab_mask_batched.side_effect = _serial_batched_fill
         active.allocate_vocab_mask.return_value = torch.zeros(3, VOCAB_SIZE)
         active.move_vocab_mask.return_value = torch.zeros(3, VOCAB_SIZE)
 
@@ -268,10 +276,34 @@ class TestUpdateRegexVocabMask(CustomTestCase):
         info.grammars = [active, finished, terminated]
         info.update_regex_vocab_mask()
 
+        active.fill_vocab_mask_batched.assert_called_once()
         active.fill_vocab_mask.assert_called_once()
         finished.fill_vocab_mask.assert_not_called()
         terminated.fill_vocab_mask.assert_not_called()
 
+    def test_batched_fill_skips_serial_loop(self):
+        """A backend with a batched kernel must not also run the serial fill."""
+        active = MagicMock()
+        active.finished = False
+        active.is_terminated.return_value = False
+        active.allocate_vocab_mask.return_value = torch.zeros(2, VOCAB_SIZE)
+        active.move_vocab_mask.return_value = torch.zeros(2, VOCAB_SIZE)
+
+        other = MagicMock()
+        other.finished = False
+        other.is_terminated.return_value = False
+
+        info = _make_info(batch_size=2)
+        info.grammars = [active, other]
+        info.update_regex_vocab_mask()
+
+        # The batched call happened with both active rows; serial fill skipped.
+        active.fill_vocab_mask_batched.assert_called_once()
+        entries, _mask = active.fill_vocab_mask_batched.call_args.args
+        self.assertEqual([e.row for e in entries], [0, 1])
+        active.fill_vocab_mask.assert_not_called()
+        other.fill_vocab_mask.assert_not_called()
+
 
 # filter_batch
 class TestFilterBatch(CustomTestCase):