From 9989077f24d376dfe7167bf8708c9899771bb3ce Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sat, 25 Jul 2026 16:36:32 -0700 Subject: [PATCH] Use native batched llguidance mask generation (#32412) Co-authored-by: Alec S <10566873+alecsolder@users.noreply.github.com> --- .../srt/constrained/base_grammar_backend.py | 79 +++++++++++- .../srt/constrained/llguidance_backend.py | 113 ++++++++++++++++-- .../srt/sampling/sampling_batch_info.py | 14 ++- .../test_llguidance_batched_mask.py | 112 +++++++++++++++++ .../unit/sampling/test_sampling_batch_info.py | 32 +++++ 5 files changed, 335 insertions(+), 15 deletions(-) create mode 100644 test/registered/unit/constrained/test_llguidance_batched_mask.py diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py index 76e5b52d8..494a3303c 100644 --- a/python/sglang/srt/constrained/base_grammar_backend.py +++ b/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): diff --git a/python/sglang/srt/constrained/llguidance_backend.py b/python/sglang/srt/constrained/llguidance_backend.py index b74563752..dbc223c62 100644 --- a/python/sglang/srt/constrained/llguidance_backend.py +++ b/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: diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 9a2dc6adc..b6005c571 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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) diff --git a/test/registered/unit/constrained/test_llguidance_batched_mask.py b/test/registered/unit/constrained/test_llguidance_batched_mask.py new file mode 100644 index 000000000..23b4362ee --- /dev/null +++ b/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() diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 37add9f0e..66159b2fb 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/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):