From 3da1071d56d14f2e56c000cd1490036d0a6a2167 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 25 Jul 2026 15:27:43 -0700 Subject: [PATCH] [Spec] Hold the grammar bitmask in one `GrammarMask` type across all decode paths (#32409) --- .../srt/constrained/base_grammar_backend.py | 15 +++++++++- .../sglang/srt/model_executor/model_runner.py | 4 +-- .../srt/sampling/sampling_batch_info.py | 28 ++++++++--------- python/sglang/srt/speculative/dflash_info.py | 4 --- .../sglang/srt/speculative/dflash_info_v2.py | 4 --- python/sglang/srt/speculative/dflash_utils.py | 4 +-- .../srt/speculative/dflash_worker_v2.py | 17 ++++------- .../dspark_components/dspark_verify.py | 2 +- .../dspark_components/dspark_worker_v2.py | 14 +++------ python/sglang/srt/speculative/eagle_info.py | 2 -- python/sglang/srt/speculative/eagle_utils.py | 10 +++---- .../srt/speculative/eagle_worker_common.py | 14 +++------ python/sglang/srt/speculative/ngram_info.py | 3 -- python/sglang/srt/speculative/ngram_worker.py | 9 +++--- python/sglang/srt/speculative/spec_utils.py | 30 +++++++++---------- .../unit/sampling/test_sampling_batch_info.py | 21 ++++++------- 16 files changed, 80 insertions(+), 101 deletions(-) diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py index 03ad0b450..76e5b52d8 100644 --- a/python/sglang/srt/constrained/base_grammar_backend.py +++ b/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.""" diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e6f0460e5..c919cd394 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 0d8d8c59b..9a2dc6adc 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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 - apply_mask_func: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None + grammar_mask: Optional[GrammarMask] = 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.apply_mask_func = None + self.grammar_mask = 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: - self.apply_mask_func(logits=logits, vocab_mask=self.vocab_mask) + if self.grammar_mask is not None: + self.grammar_mask.apply(logits) if self.logit_bias is not None: logits.add_(self.logit_bias) diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 6758878b5..8ec5420bc 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/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: diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 29b593fd6..72162ba51 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/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. diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 246b7a1ca..7df8edbd0 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/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, diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 58a3a11d5..0be95e06d 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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 - # 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( + grammar_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: - verify_input.grammar.apply_vocab_mask( - logits=logits_output.next_token_logits, vocab_mask=vocab_mask - ) + if grammar_mask is not None: + grammar_mask.apply(logits_output.next_token_logits) candidates = draft_tokens new_seq_lens = None diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 9e879862c..4d5478cd5 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/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 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 2ac5d2b2c..ae778ee1d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/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: - draft_input.grammar.apply_vocab_mask( - logits=logits_output.next_token_logits, vocab_mask=vocab_mask - ) + if grammar_mask is not None: + grammar_mask.apply(logits_output.next_token_logits) epilogue = self._verify_executor.verify_epilogue folded_accept = fold_eligible and run_compact and can_run_cuda_graph diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index e7b724165..4184fe465 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 97592cc99..780261541 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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: - assert verify_input.grammar is not None - verify_input.grammar.apply_vocab_mask( - logits=next_token_logits, vocab_mask=vocab_mask - ) + if grammar_mask is not None: + grammar_mask.apply(next_token_logits) candidates = verify_input.draft_token.reshape(bs, verify_input.draft_token_num) predict_shape = list(next_token_logits.shape)[:-1] diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 69498b66b..56e645d4e 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/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 - # 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( + grammar_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(), diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 11d63fb9e..50d8e0945 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index d2616acd6..38709a34a 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/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, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index c05d2dfc5..7eb35b982 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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 + return allocate_token_bitmask, grammar 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 - both host work, so both overlap that forward. + Call it after the target verify launch -- every step here is host work, so it all + 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 - return vocab_mask + sampling_info.grammar_mask = None + return GrammarMask(grammar, vocab_mask) def load_token_map(token_map_path: str) -> List[int]: diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 0ff0efbd8..37add9f0e 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/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) - info.apply_mask_func = MagicMock() + grammar = 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.apply_mask_func) + self.assertIsNone(info.grammar_mask) 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."""