From b58614126fd46e909ba92101e9435c64a7df9835 Mon Sep 17 00:00:00 2001 From: Doug Brown Date: Thu, 17 Sep 2026 17:44:43 -0400 Subject: [PATCH 1/2] =?UTF-8?q?=F0=9F=9B=A1=EF=B8=8F=20Refuse=20to=20gener?= =?UTF-8?q?ate=20unconstrained=20when=20a=20grammar=20fails=20to=20compile?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A json_schema, regex_pattern, or grammar_string that LLGuidance rejects was silently dropped: the filter was skipped, generation ran unconstrained, and the client received a 200 with output that ignores the requested constraint. The detailed error went to server stdout only. Fail closed instead. Filter construction now raises GrammarParseError, a new client-error type in common/errors.py alongside ContextLengthExceededError: - OAI chat/completions and Kobold non-streaming handlers map it to HTTPException(400) via the existing handle_request_error path; streaming handlers yield an error event carrying the compile failure. - A missing or broken llguidance backend re-raises the original exception so it keeps mapping to a server error (503), not a client error. - The per-request schema/grammar is logged once with its compile exception before the rejection, replacing the traceback.print_exc() stdout dump. --- backends/exllamav3/grammar.py | 52 +++++-- common/errors.py | 8 + endpoints/Kobold/utils/generation.py | 11 +- endpoints/OAI/utils/chat_completion.py | 13 +- endpoints/OAI/utils/completion.py | 13 +- tests/test_grammar_fail_closed.py | 196 +++++++++++++++++++++++++ 6 files changed, 275 insertions(+), 18 deletions(-) create mode 100644 tests/test_grammar_fail_closed.py diff --git a/backends/exllamav3/grammar.py b/backends/exllamav3/grammar.py index 57e9197a..ee481313 100644 --- a/backends/exllamav3/grammar.py +++ b/backends/exllamav3/grammar.py @@ -6,9 +6,26 @@ Filter, LLGuidanceFilter, ) +from common.errors import GrammarParseError from common.logger import xlogger +def _llguidance_ready() -> bool: + """Whether the llguidance backend itself is importable and present. + + A missing backend is a server environment problem and keeps its own + exception identity (server error), while a schema/regex/grammar the + backend rejects is a client error. + """ + + try: + from exllamav3.generator.filter.llguidance import llguidance_available + except ImportError: + # Flag moved in a future exllamav3; let construction errors decide. + return True + return llguidance_available + + class ExLlamaV3Grammar: """ExLlamaV3 class for various grammar filters/parsers.""" @@ -36,14 +53,15 @@ def add_json_schema_filter( json_schema=schema, trigger_token=trigger_token_id, ) - except Exception: - traceback.print_exc() + except Exception as exc: xlogger.error( - "Skipping because the JSON schema couldn't be parsed. " - "Please read the above error for more information.", + "JSON schema could not be compiled; rejecting the request instead " + "of generating without the constraint.", {"schema": schema, "exception": traceback.format_exc()}, ) - return + if not _llguidance_ready(): + raise + raise GrammarParseError(f"The JSON schema could not be compiled: {exc}") from exc self.filters.append(lmfilter) @@ -62,14 +80,15 @@ def add_regex_filter( regex=pattern, trigger_token=trigger_token_id, ) - except Exception: - traceback.print_exc() + except Exception as exc: xlogger.error( - "Skipping because the regex pattern couldn't be parsed. " - "Please read the above error for more information.", + "Regex pattern could not be compiled; rejecting the request instead " + "of generating without the constraint.", {"pattern": pattern, "exception": traceback.format_exc()}, ) - return + if not _llguidance_ready(): + raise + raise GrammarParseError(f"The regex pattern could not be compiled: {exc}") from exc self.filters.append(lmfilter) @@ -94,13 +113,16 @@ def add_grammar_filter( trigger_token=trigger_token_id, **{grammar_kind: grammar_string}, ) - except Exception: - traceback.print_exc() + except Exception as exc: xlogger.error( - "Skipping because the grammar couldn't be parsed. " - "Please read the above error for more information.", + f"{grammar_kind} could not be compiled; rejecting the request instead " + "of generating without the constraint.", {"grammar_string": grammar_string, "exception": traceback.format_exc()}, ) - return + if not _llguidance_ready(): + raise + raise GrammarParseError( + f"The grammar ({grammar_kind}) could not be compiled: {exc}" + ) from exc self.filters.append(lmfilter) diff --git a/common/errors.py b/common/errors.py index a4b098ae..93ffbc7a 100644 --- a/common/errors.py +++ b/common/errors.py @@ -6,6 +6,14 @@ class ContextLengthExceededError(ValueError): """Raised when a tokenized prompt exceeds the loaded model's context limit.""" +class GrammarParseError(ValueError): + """Raised when a requested json_schema, regex, or grammar fails to compile. + + Constrained decoding refuses to fall back to unconstrained generation, so + the request is rejected instead of silently served without the constraint. + """ + + class ContextLengthHTTPException(HTTPException): """HTTP error for OpenAI-compatible context overflow responses.""" diff --git a/endpoints/Kobold/utils/generation.py b/endpoints/Kobold/utils/generation.py index db59b88d..1963a1d7 100644 --- a/endpoints/Kobold/utils/generation.py +++ b/endpoints/Kobold/utils/generation.py @@ -5,7 +5,11 @@ from sse_starlette.event import ServerSentEvent from common import model -from common.errors import ContextLengthExceededError, ContextLengthHTTPException +from common.errors import ( + ContextLengthExceededError, + ContextLengthHTTPException, + GrammarParseError, +) from common.networking import ( get_context_length_generator_error, get_generator_error, @@ -102,6 +106,8 @@ async def stream_generation(data: GenerateRequest, request: Request): yield ServerSentEvent(event="message", data=response.model_dump_json(), sep="\n") except ContextLengthExceededError as exc: yield get_context_length_generator_error(str(exc)) + except GrammarParseError as exc: + yield get_generator_error(str(exc), exc_info=False) except Exception: yield get_generator_error( f"{request_tag(request)} kobold/generate aborted. Please check the server console." @@ -125,6 +131,9 @@ async def get_generation(data: GenerateRequest, request: Request): except ContextLengthExceededError as exc: error_message = handle_request_error(str(exc), exc_info=False).error.message raise ContextLengthHTTPException(error_message) from exc + except GrammarParseError as exc: + error_message = handle_request_error(str(exc), exc_info=False).error.message + raise HTTPException(400, error_message) from exc except Exception as exc: error_message = handle_request_error( f"{request_tag(request)} kobold/generate aborted. Maybe the model was unloaded? " diff --git a/endpoints/OAI/utils/chat_completion.py b/endpoints/OAI/utils/chat_completion.py index 5c2b350b..c55a24b8 100644 --- a/endpoints/OAI/utils/chat_completion.py +++ b/endpoints/OAI/utils/chat_completion.py @@ -8,7 +8,11 @@ from typing import List, Optional from fastapi import HTTPException, Request from jinja2 import TemplateError -from common.errors import ContextLengthExceededError, ContextLengthHTTPException +from common.errors import ( + ContextLengthExceededError, + ContextLengthHTTPException, + GrammarParseError, +) from common.logger import xlogger import re @@ -912,6 +916,9 @@ async def stream_generate_chat_completion( except ContextLengthExceededError as exc: yield get_context_length_generator_error(str(exc)) + except GrammarParseError as exc: + yield get_generator_error(str(exc), exc_info=False) + except Exception as e: xlogger.error("Error during chat completion", str(e), details=f"\n{str(e)}") yield get_generator_error("Chat completion aborted. Please check the server console.") @@ -986,6 +993,10 @@ async def generate_chat_completion( error_message = handle_request_error(str(exc), exc_info=False).error.message raise ContextLengthHTTPException(error_message) from exc + except GrammarParseError as exc: + error_message = handle_request_error(str(exc), exc_info=False).error.message + raise HTTPException(400, error_message) from exc + except Exception as exc: error_message = handle_request_error( f"{request_tag(request)} chat completion aborted. Maybe the model was unloaded? " diff --git a/endpoints/OAI/utils/completion.py b/endpoints/OAI/utils/completion.py index e1f127b1..0d047172 100644 --- a/endpoints/OAI/utils/completion.py +++ b/endpoints/OAI/utils/completion.py @@ -11,7 +11,11 @@ from time import time from fastapi import HTTPException, Request -from common.errors import ContextLengthExceededError, ContextLengthHTTPException +from common.errors import ( + ContextLengthExceededError, + ContextLengthHTTPException, + GrammarParseError, +) from common.logger import xlogger from typing import List, Optional @@ -353,6 +357,9 @@ async def stream_generate_completion( except ContextLengthExceededError as exc: yield get_context_length_generator_error(str(exc)) + except GrammarParseError as exc: + yield get_generator_error(str(exc), exc_info=False) + except Exception as e: xlogger.error("Error during completion", str(e), details=f"\n{str(e)}") yield get_generator_error("Completion aborted. Please check the server console.") @@ -431,6 +438,10 @@ async def generate_completion( error_message = handle_request_error(str(exc), exc_info=False).error.message raise ContextLengthHTTPException(error_message) from exc + except GrammarParseError as exc: + error_message = handle_request_error(str(exc), exc_info=False).error.message + raise HTTPException(400, error_message) from exc + except Exception as exc: error_message = handle_request_error( f"{request_tag(request)} completion aborted. Maybe the model was unloaded? " diff --git a/tests/test_grammar_fail_closed.py b/tests/test_grammar_fail_closed.py new file mode 100644 index 00000000..983814f0 --- /dev/null +++ b/tests/test_grammar_fail_closed.py @@ -0,0 +1,196 @@ +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +from fastapi import HTTPException + +from backends.exllamav3.grammar import ExLlamaV3Grammar +from common import model +from common.errors import GrammarParseError +from endpoints.Kobold.utils import generation as kobold_generation +from endpoints.OAI.utils import chat_completion, completion + + +class DummyDisconnectHandler: + async def cleanup(self): + pass + + +class DummyRequestData: + n = 1 + stream_options = None + + def model_copy(self, deep=False): + return self + + def model_dump(self, mode=None): + return {} + + +def request_with_id(request_id="request-id"): + return SimpleNamespace(state=SimpleNamespace(id=request_id)) + + +def _raising_filter(message="Invalid grammar: boom"): + def factory(*args, **kwargs): + raise ValueError(message) + + return factory + + +class GrammarParseErrorUnitTests(unittest.TestCase): + """A rejected schema/regex/grammar must fail the request, not silently + degrade to unconstrained generation.""" + + def setUp(self): + self.handler = ExLlamaV3Grammar() + self.tokenizer = SimpleNamespace( + actual_vocab_size=4, + extended_id_to_piece={0: "a", 1: "b", 2: " ", 3: "\n"}, + ) + + def test_bad_json_schema_raises(self): + with patch("backends.exllamav3.grammar.LLGuidanceFilter", _raising_filter()): + with self.assertRaises(GrammarParseError): + self.handler.add_json_schema_filter({"type": "object"}, self.tokenizer) + self.assertEqual(self.handler.filters, []) + + def test_bad_regex_raises(self): + with patch("backends.exllamav3.grammar.LLGuidanceFilter", _raising_filter()): + with self.assertRaises(GrammarParseError): + self.handler.add_regex_filter("[unclosed", self.tokenizer) + self.assertEqual(self.handler.filters, []) + + def test_bad_grammar_raises(self): + with patch("backends.exllamav3.grammar.LLGuidanceFilter", _raising_filter()): + with self.assertRaises(GrammarParseError): + self.handler.add_grammar_filter("rule ::= ", self.tokenizer) + self.assertEqual(self.handler.filters, []) + + def test_valid_schema_appends_filter(self): + sentinel = SimpleNamespace( + name="filter", + tokenizer=None, + trigger_token=None, + prefix_str=None, + eos_after_completed=True, + ) + with patch( + "backends.exllamav3.grammar.LLGuidanceFilter", + lambda *args, **kwargs: sentinel, + ): + self.handler.add_json_schema_filter({"type": "object"}, self.tokenizer) + self.assertEqual(self.handler.filters, [sentinel]) + + def test_missing_llguidance_dependency_is_not_masked(self): + # A server-side environment problem must keep its original identity + # (it maps to a server error), not a client-facing parse error. + with patch("backends.exllamav3.grammar.LLGuidanceFilter", _raising_filter()): + with patch( + "backends.exllamav3.grammar._llguidance_ready", return_value=False + ): + with self.assertRaises(ValueError) as raised: + self.handler.add_json_schema_filter({"type": "object"}, self.tokenizer) + self.assertNotIsInstance(raised.exception, GrammarParseError) + + +class GrammarParseErrorEndpointTests(unittest.IsolatedAsyncioTestCase): + """Grammar parse errors are client errors (400), unlike generation faults.""" + + async def test_chat_completion_returns_400_for_grammar_parse_error(self): + async def collector(*args, **kwargs): + raise GrammarParseError( + "The JSON schema could not be compiled: Invalid grammar: boom" + ) + + original_container = model.container + model.container = SimpleNamespace(reasoning=False, harmony=False, muse_glimmer=False) + try: + with patch.object(chat_completion, "_chat_stream_collector", collector): + with self.assertRaises(HTTPException) as raised: + await chat_completion.generate_chat_completion( + "prompt", + None, + DummyRequestData(), + request_with_id(), + Path("model"), + DummyDisconnectHandler(), + ) + finally: + model.container = original_container + + self.assertEqual(raised.exception.status_code, 400) + self.assertIn("could not be compiled", raised.exception.detail) + + async def test_completion_returns_400_for_grammar_parse_error(self): + async def collector(*args, **kwargs): + raise GrammarParseError( + "The JSON schema could not be compiled: Invalid grammar: boom" + ) + + with patch.object(completion, "_stream_collector", collector): + with self.assertRaises(HTTPException) as raised: + await completion.generate_completion( + "prompt", + DummyRequestData(), + request_with_id(), + Path("model"), + DummyDisconnectHandler(), + ) + + self.assertEqual(raised.exception.status_code, 400) + self.assertIn("could not be compiled", raised.exception.detail) + + async def test_kobold_generation_returns_400_for_grammar_parse_error(self): + async def collector(*args, **kwargs): + raise GrammarParseError("The JSON schema could not be compiled: Invalid grammar: boom") + yield + + data = SimpleNamespace(genkey=None) + with patch.object(kobold_generation, "_stream_collector", collector): + with self.assertRaises(HTTPException) as raised: + await kobold_generation.get_generation(data, request_with_id()) + + self.assertEqual(raised.exception.status_code, 400) + self.assertIn("could not be compiled", raised.exception.detail) + + +class GrammarParseErrorStreamTests(unittest.IsolatedAsyncioTestCase): + """Streaming requests surface the grammar rejection as an error event.""" + + async def test_chat_stream_yields_error_event_for_grammar_parse_error(self): + error = GrammarParseError( + "The JSON schema could not be compiled: Invalid grammar: boom" + ) + + async def collector(task_idx, gen_queue=None, *args, **kwargs): + # Mirror the real collector's contract: failures are pushed to the + # queue as exception objects and re-raised by the consumer loop. + await gen_queue.put(error) + + original_container = model.container + model.container = SimpleNamespace(reasoning=False, harmony=False, muse_glimmer=False) + try: + with patch.object(chat_completion, "_chat_stream_collector", collector): + chunks = [ + chunk + async for chunk in chat_completion.stream_generate_chat_completion( + "prompt", + None, + DummyRequestData(), + request_with_id(), + Path("model"), + DummyDisconnectHandler(), + ) + ] + finally: + model.container = original_container + + joined = "".join(chunks) + self.assertIn("could not be compiled", joined) + self.assertNotIn('"content"', joined) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file From e500f97f446e38e3683b073d3be88be5043972b6 Mon Sep 17 00:00:00 2001 From: Doug Brown Date: Thu, 17 Sep 2026 17:44:50 -0400 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=9B=A4=EF=B8=8F=20Re-rail=20grammar-c?= =?UTF-8?q?onstrained=20generation=20out=20of=20whitespace=20stalls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A grammar filter can only forbid tokens; it cannot compel the model to continue or stop. When the legal continuation set includes unlimited whitespace, a model that "feels done" (e.g. it skipped a trailing required property) can emit whitespace forever: every sampled token is legal, EOS and closing braces are masked, and the job burns its whole token budget on output that never satisfies the grammar. ExLlamaV3's loop detector does not catch this - its default 800-token window is tuned for long repeating sequences, not short legal attractors. WhitespaceStallGuardFilter wraps each grammar filter in Tabby: - Tracks accepted tokens. After stall_tokens (8) consecutive pure-whitespace tokens while the inner grammar is still incomplete, the next logit mask ANDs out whitespace, forcing the sampler onto a legal content token (or EOS). Each re-rail advances the grammar state, so repeated stalls terminate. - If excluding whitespace would empty the allowed set (the grammar truly requires whitespace there), the unmasked grammar mask is returned. - Delegates feed/rewind/rollback/reset/attach and mirrors is_active, trigger_token and eos_after_completed into the wrapped filter so job-level activation and deactivation keep working; pre-trigger history is cleared on activation. - Whitespace token ids come from one packed bitmask built lazily per handler from the tokenizer vocab (inner get_vocab, extended_id_to_piece fallback; SentencePiece "▁" counts as space; special tokens never count). If the bitmask cannot be built the guard disables itself and generation behaves as before. Grammar handlers wrap every constructed LLGuidanceFilter with a request-level threshold (stall_guard_tokens, default 8, 0/null disables, logged with the request settings); no exllamav3 changes are required. --- backends/exllamav3/grammar.py | 36 +++- backends/exllamav3/model.py | 4 +- backends/exllamav3/whitespace_guard.py | 177 ++++++++++++++++++++ common/sampling.py | 13 ++ tests/test_grammar_fail_closed.py | 42 ++++- tests/test_whitespace_stall_guard.py | 220 +++++++++++++++++++++++++ 6 files changed, 485 insertions(+), 7 deletions(-) create mode 100644 backends/exllamav3/whitespace_guard.py create mode 100644 tests/test_whitespace_stall_guard.py diff --git a/backends/exllamav3/grammar.py b/backends/exllamav3/grammar.py index ee481313..3dde62d5 100644 --- a/backends/exllamav3/grammar.py +++ b/backends/exllamav3/grammar.py @@ -6,6 +6,10 @@ Filter, LLGuidanceFilter, ) +from backends.exllamav3.whitespace_guard import ( + WhitespaceStallGuardFilter, + build_whitespace_bitmask, +) from common.errors import GrammarParseError from common.logger import xlogger @@ -31,8 +35,32 @@ class ExLlamaV3Grammar: filters: List[Filter] - def __init__(self): + def __init__(self, stall_tokens: Optional[int] = 8): self.filters = [] + self.stall_tokens = stall_tokens + self._whitespace_bitmask = None + + def _wrap_with_stall_guard(self, lmfilter: Filter, tokenizer: Tokenizer) -> Filter: + """Wrap the filter with the whitespace stall guard, building the + vocab-wide whitespace bitmask once per handler.""" + + if not self.stall_tokens or self.stall_tokens < 1: + return lmfilter + if self._whitespace_bitmask is None: + try: + self._whitespace_bitmask = build_whitespace_bitmask(tokenizer) + except Exception as exc: + xlogger.warning( + "Whitespace stall guard unavailable; constrained generation " + "may stall in legal whitespace runs.", + {"exception": str(exc)}, + ) + self._whitespace_bitmask = False + if self._whitespace_bitmask is False: + return lmfilter + return WhitespaceStallGuardFilter( + lmfilter, self._whitespace_bitmask, stall_tokens=self.stall_tokens + ) def add_json_schema_filter( self, @@ -63,7 +91,7 @@ def add_json_schema_filter( raise raise GrammarParseError(f"The JSON schema could not be compiled: {exc}") from exc - self.filters.append(lmfilter) + self.filters.append(self._wrap_with_stall_guard(lmfilter, tokenizer)) def add_regex_filter( self, @@ -90,7 +118,7 @@ def add_regex_filter( raise raise GrammarParseError(f"The regex pattern could not be compiled: {exc}") from exc - self.filters.append(lmfilter) + self.filters.append(self._wrap_with_stall_guard(lmfilter, tokenizer)) def add_grammar_filter( self, @@ -125,4 +153,4 @@ def add_grammar_filter( f"The grammar ({grammar_kind}) could not be compiled: {exc}" ) from exc - self.filters.append(lmfilter) + self.filters.append(self._wrap_with_stall_guard(lmfilter, tokenizer)) diff --git a/backends/exllamav3/model.py b/backends/exllamav3/model.py index 57e750bb..13718bb2 100644 --- a/backends/exllamav3/model.py +++ b/backends/exllamav3/model.py @@ -1360,7 +1360,7 @@ async def generate_gen( prompts = [prompt] stop_conditions = params.stop add_bos_token = unwrap(params.add_bos_token, self.hf_model.add_bos_token()) - grammar_handler = ExLlamaV3Grammar() + grammar_handler = ExLlamaV3Grammar(stall_tokens=params.stall_guard_tokens) # Get multimodal embeddings if present mm_embeddings_content = mm_embeddings.content if mm_embeddings else [] @@ -1460,6 +1460,8 @@ async def generate_gen( settings.append(("regex_pattern", True)) if params.grammar_string: settings.append(("grammar_string", True)) + if params.json_schema or params.regex_pattern or params.grammar_string: + settings.append(("stall_guard_tokens", params.stall_guard_tokens)) if params.token_healing: settings.append(("token_healing", True)) if params.logprobs or params.top_logprobs: diff --git a/backends/exllamav3/whitespace_guard.py b/backends/exllamav3/whitespace_guard.py new file mode 100644 index 00000000..ad9834ed --- /dev/null +++ b/backends/exllamav3/whitespace_guard.py @@ -0,0 +1,177 @@ +from collections import deque + +import torch + +from exllamav3 import Filter + +from common.logger import xlogger + + +def build_whitespace_bitmask(tokenizer) -> torch.Tensor: + """Packed int32 bitmask (bit set = token is pure whitespace) over the vocab. + + A token counts as whitespace when its piece is empty or contains only + whitespace, counting SentencePiece's "▁" as a space. Special tokens are + never treated as whitespace. + """ + + vocab_size = tokenizer.actual_vocab_size + piece_map = None + inner = getattr(tokenizer, "tokenizer", None) + if inner is not None and hasattr(inner, "get_vocab"): + try: + piece_map = { + token_id: piece + for piece, token_id in inner.get_vocab().items() + if token_id < vocab_size + } + except Exception: + piece_map = None + if piece_map is None: + piece_map = getattr(tokenizer, "extended_id_to_piece", None) or {} + words = (vocab_size + 31) // 32 + mask = torch.zeros(words, dtype=torch.int32) + for token_id, piece in piece_map.items(): + if piece is None or piece.startswith("<") and piece.endswith(">"): + continue + flat = piece.replace("▁", " ") + if flat != "" and flat.strip() == "": + mask[token_id >> 5] |= 1 << (token_id & 31) + return mask + + +class WhitespaceStallGuardFilter(Filter): + """Re-rails a grammar-constrained generation out of whitespace stalls. + + A grammar can only forbid tokens; it cannot compel the model to continue + or stop. When the remaining legal continuations include unlimited + whitespace, a model that "feels done" can emit whitespace forever, + burning the token budget with output that never satisfies the grammar. + + This wrapper watches accepted tokens. If the model emits `stall_tokens` + consecutive pure-whitespace tokens while the inner grammar filter is + still incomplete, the next logit mask additionally excludes whitespace + tokens, forcing the sampler onto a legal content token (or EOS). This + advances the grammar instead of idling in a legal whitespace attractor. + + If excluding whitespace would empty the allowed set (the grammar truly + requires whitespace), the unmasked grammar mask is returned unchanged. + """ + + def __init__(self, inner: Filter, whitespace_bitmask: torch.Tensor, stall_tokens: int = 8): + if stall_tokens < 1: + raise ValueError("stall_tokens must be at least 1") + + # Assign inner before super().__init__: the base class's is_active + # assignment flows through the property setter, which mirrors into + # the inner filter. + self.inner = inner + self._is_active = False + self._recent: deque[int] = deque(maxlen=stall_tokens) + self._stalled = False + super().__init__( + inner.tokenizer, + inner.trigger_token, + inner.prefix_str, + inner.eos_after_completed, + ) + self.whitespace_bitmask = whitespace_bitmask + self.stall_tokens = stall_tokens + # Per-token boolean view of the packed bitmask, for dense masks. + bits = ( + (whitespace_bitmask.unsqueeze(-1) >> torch.arange(32, dtype=torch.int32)) & 1 + ).bool().reshape(-1) + self._ws_token_bool = bits + + @property + def is_active(self): + return self._is_active + + @is_active.setter + def is_active(self, value): + was = self._is_active + self._is_active = value + self.inner.is_active = value + if value and not was: + # Activating from a trigger: pre-trigger history is not part of + # the constrained sequence. + self._recent.clear() + self._stalled = False + + # -- token bookkeeping ------------------------------------------------- + + def _is_whitespace(self, token: int) -> bool: + words = self.whitespace_bitmask + return bool(token >= 0 and (words[token >> 5].item() >> (token & 31)) & 1) + + def _update_stall_state(self): + self._stalled = ( + len(self._recent) == self.stall_tokens + and not self.inner.is_completed() + and all(self._is_whitespace(t) for t in self._recent) + ) + + # -- Filter interface --------------------------------------------------- + + def feed(self, token: int) -> bool: + eos = self.inner.feed(token) + self._recent.append(token) + self._update_stall_state() + return eos + + def rewind(self, num_tokens: int): + self.inner.rewind(num_tokens) + for _ in range(min(num_tokens, len(self._recent))): + self._recent.pop() + self._update_stall_state() + + def rollback_tokens(self, num_tokens: int) -> bool: + eos = self.inner.rollback_tokens(num_tokens) + for _ in range(min(num_tokens, len(self._recent))): + self._recent.pop() + self._update_stall_state() + return eos + + def reset(self): + self.inner.reset() + self._recent.clear() + self._stalled = False + + def is_completed(self) -> bool: + return self.inner.is_completed() + + def use_background_worker(self) -> bool: + return self.inner.use_background_worker() + + def attach(self, job): + super().attach(job) + self.inner.attach(job) + + def get_next_logit_mask(self) -> torch.Tensor: + mask = self.inner.get_next_logit_mask() + if not self._stalled: + return mask + + if mask.dtype == torch.int32: + # Packed bitmask: bit set = allowed. Clear whitespace bits. + words = self.whitespace_bitmask + guarded = mask & ~words[: mask.shape[-1]].to(mask.dtype) + if (guarded != 0).any(): + return guarded + xlogger.debug( + "Whitespace stall guard active but every legal continuation " + "is whitespace; leaving the grammar mask unchanged." + ) + return mask + + # Dense additive mask: 0 = allowed, -inf = masked. + guarded = mask.clone() + ws = self._ws_token_bool[: guarded.shape[-1]] + guarded[..., ws] = float("-inf") + if torch.isfinite(guarded).any(): + return guarded + xlogger.debug( + "Whitespace stall guard active but every legal continuation " + "is whitespace; leaving the grammar mask unchanged." + ) + return mask \ No newline at end of file diff --git a/common/sampling.py b/common/sampling.py index 8e731a41..6ae45fb1 100644 --- a/common/sampling.py +++ b/common/sampling.py @@ -323,6 +323,19 @@ class BaseSamplerRequest(BaseModel): ge=0, ) + stall_guard_tokens: Optional[int] = Field( + default_factory=lambda: get_default_sampler_value("stall_guard_tokens", 8), + description=( + "ExLlamaV3 only. Re-rails grammar-constrained generation " + "(json_schema, regex_pattern, grammar_string) out of legal " + "whitespace runs: after N consecutive whitespace-only tokens while " + "the grammar is still incomplete, the next logit mask excludes " + "whitespace so sampling lands on a legal content token or EOS. " + "Set 0 or null to disable." + ), + ge=0, + ) + def param_source(self, name: str) -> str: """ Where the effective value of a sampler param came from: "req" (sent with the diff --git a/tests/test_grammar_fail_closed.py b/tests/test_grammar_fail_closed.py index 983814f0..239b4bc3 100644 --- a/tests/test_grammar_fail_closed.py +++ b/tests/test_grammar_fail_closed.py @@ -6,6 +6,7 @@ from fastapi import HTTPException from backends.exllamav3.grammar import ExLlamaV3Grammar +from backends.exllamav3.whitespace_guard import WhitespaceStallGuardFilter from common import model from common.errors import GrammarParseError from endpoints.Kobold.utils import generation as kobold_generation @@ -68,7 +69,41 @@ def test_bad_grammar_raises(self): self.handler.add_grammar_filter("rule ::= ", self.tokenizer) self.assertEqual(self.handler.filters, []) - def test_valid_schema_appends_filter(self): + def test_stall_guard_tokens_zero_disables_guard(self): + sentinel = SimpleNamespace( + name="filter", + tokenizer=None, + trigger_token=None, + prefix_str=None, + eos_after_completed=True, + ) + handler = ExLlamaV3Grammar(stall_tokens=0) + with patch( + "backends.exllamav3.grammar.LLGuidanceFilter", + lambda *args, **kwargs: sentinel, + ): + handler.add_json_schema_filter({"type": "object"}, self.tokenizer) + self.assertEqual(handler.filters, [sentinel]) + + def test_stall_guard_threshold_flows_to_wrapper(self): + sentinel = SimpleNamespace( + name="filter", + tokenizer=None, + trigger_token=None, + prefix_str=None, + eos_after_completed=True, + ) + handler = ExLlamaV3Grammar(stall_tokens=4) + with patch( + "backends.exllamav3.grammar.LLGuidanceFilter", + lambda *args, **kwargs: sentinel, + ): + handler.add_json_schema_filter({"type": "object"}, self.tokenizer) + self.assertEqual(len(handler.filters), 1) + self.assertIsInstance(handler.filters[0], WhitespaceStallGuardFilter) + self.assertEqual(handler.filters[0].stall_tokens, 4) + + def test_valid_schema_appends_guarded_filter(self): sentinel = SimpleNamespace( name="filter", tokenizer=None, @@ -81,7 +116,10 @@ def test_valid_schema_appends_filter(self): lambda *args, **kwargs: sentinel, ): self.handler.add_json_schema_filter({"type": "object"}, self.tokenizer) - self.assertEqual(self.handler.filters, [sentinel]) + self.assertEqual(len(self.handler.filters), 1) + wrapped = self.handler.filters[0] + self.assertIsInstance(wrapped, WhitespaceStallGuardFilter) + self.assertIs(wrapped.inner, sentinel) def test_missing_llguidance_dependency_is_not_masked(self): # A server-side environment problem must keep its original identity diff --git a/tests/test_whitespace_stall_guard.py b/tests/test_whitespace_stall_guard.py new file mode 100644 index 00000000..8afc6f3f --- /dev/null +++ b/tests/test_whitespace_stall_guard.py @@ -0,0 +1,220 @@ +import unittest +from types import SimpleNamespace + +import torch + +from backends.exllamav3.whitespace_guard import ( + WhitespaceStallGuardFilter, + build_whitespace_bitmask, +) + + +VOCAB = 8 +PIECES = { + 0: "", + 1: "hello", + 2: " ", + 3: "\n", + 4: "▁", + 5: "▁▁", + 6: ",", + 7: "x", +} +WS_IDS = {2, 3, 4, 5} + + +def fake_tokenizer(): + return SimpleNamespace(actual_vocab_size=VOCAB, extended_id_to_piece=dict(PIECES)) + + +def packed_mask(allowed_ids): + mask = torch.zeros((VOCAB + 31) // 32, dtype=torch.int32) + for token_id in allowed_ids: + mask[token_id >> 5] |= 1 << (token_id & 31) + return mask + + +def dense_mask(allowed_ids): + mask = torch.zeros((1, VOCAB), dtype=torch.half) + for token_id in range(VOCAB): + if token_id not in allowed_ids: + mask[0, token_id] = float("-inf") + return mask + + +class FakeInner: + def __init__(self, mask, trigger_token=None): + self.tokenizer = None + self.trigger_token = trigger_token + self.prefix_str = None + self.eos_after_completed = True + self.is_active = trigger_token is None + self._mask = mask + self._completed = False + self._eos_after_feed = False + self.fed = [] + self.rewound = [] + self.rolled_back = [] + self.attached = None + self.reset_count = 0 + + def feed(self, token): + self.fed.append(token) + return self._eos_after_feed + + def rewind(self, num_tokens): + self.rewound.append(num_tokens) + + def rollback_tokens(self, num_tokens): + self.rolled_back.append(num_tokens) + return False + + def reset(self): + self.reset_count += 1 + + def is_completed(self): + return self._completed + + def use_background_worker(self): + return True + + def attach(self, job): + self.attached = job + + def get_next_logit_mask(self): + return self._mask + + +class BuildBitmaskTests(unittest.TestCase): + def test_whitespace_pieces_flagged(self): + mask = build_whitespace_bitmask(fake_tokenizer()) + flagged = {i for i in range(VOCAB) if (mask[i >> 5].item() >> (i & 31)) & 1} + self.assertEqual(flagged, WS_IDS) + + def test_missing_piece_map_is_safe(self): + tok = SimpleNamespace(actual_vocab_size=4, extended_id_to_piece={}) + mask = build_whitespace_bitmask(tok) + self.assertEqual(mask.abs().sum().item(), 0) + + +class StallGuardTests(unittest.TestCase): + def guard(self, inner, stall_tokens=8): + return WhitespaceStallGuardFilter(inner, build_whitespace_bitmask(fake_tokenizer()), stall_tokens) + + def test_inactive_below_threshold(self): + inner = FakeInner(packed_mask({1, 2, 6})) + guard = self.guard(inner) + for token in [2, 3, 4, 5, 2, 3, 4]: # 7 whitespace tokens + guard.feed(token) + self.assertFalse(guard._stalled) + self.assertTrue(torch.equal(guard.get_next_logit_mask(), inner._mask)) + + def test_activates_at_threshold_and_clears_whitespace(self): + inner = FakeInner(packed_mask({1, 2, 3, 6})) + guard = self.guard(inner) + for token in [2, 3, 4, 5, 2, 3, 4, 3]: # 8 whitespace tokens + guard.feed(token) + self.assertTrue(guard._stalled) + mask = guard.get_next_logit_mask() + allowed = {i for i in range(VOCAB) if (mask[i >> 5].item() >> (i & 31)) & 1} + self.assertEqual(allowed, {1, 6}) # whitespace (2,3) railed out + + def test_non_whitespace_token_resets_run(self): + inner = FakeInner(packed_mask({1, 2, 6})) + guard = self.guard(inner) + for token in [2, 3, 4, 5, 2, 3, 4, 1, 2, 3]: + guard.feed(token) + self.assertFalse(guard._stalled) + + def test_never_stalls_when_inner_completed(self): + inner = FakeInner(packed_mask({1, 2, 6})) + inner._completed = True + guard = self.guard(inner) + for token in WS_IDS: + guard.feed(token) + for token in WS_IDS: + guard.feed(token) + self.assertFalse(guard._stalled) + + def test_leaves_mask_when_all_legal_tokens_are_whitespace(self): + inner = FakeInner(packed_mask({2, 3})) + guard = self.guard(inner) + for token in WS_IDS: + guard.feed(token) + for token in WS_IDS: + guard.feed(token) + self.assertTrue(guard._stalled) + self.assertTrue(torch.equal(guard.get_next_logit_mask(), inner._mask)) + + def test_rewind_retracts_stall(self): + inner = FakeInner(packed_mask({1, 2, 6})) + guard = self.guard(inner) + for token in [*WS_IDS, *WS_IDS]: # 8 ws tokens + guard.feed(token) + guard.rewind(3) # back to 5 ws + self.assertEqual(inner.rewound, [3]) + self.assertFalse(guard._stalled) + for token in [4, 5, 2]: # back to 8 ws + guard.feed(token) + self.assertTrue(guard._stalled) + + def test_rollback_tokens_delegates_and_retracts(self): + inner = FakeInner(packed_mask({1, 2, 6})) + guard = self.guard(inner) + for token in [*WS_IDS, *WS_IDS]: # 8 ws tokens + guard.feed(token) + self.assertTrue(guard._stalled) + self.assertFalse(guard.rollback_tokens(2)) + self.assertEqual(inner.rolled_back, [2]) + self.assertFalse(guard._stalled) + + def test_feed_eos_condition_propagates(self): + inner = FakeInner(packed_mask({1, 2, 6})) + inner._eos_after_feed = True + guard = self.guard(inner) + self.assertTrue(guard.feed(2)) + inner._eos_after_feed = False + self.assertFalse(guard.feed(3)) + + def test_is_active_mirrors_into_inner(self): + inner = FakeInner(packed_mask({1, 2, 6}), trigger_token=1) + guard = self.guard(inner) + self.assertFalse(guard.is_active) + self.assertFalse(inner.is_active) + guard.is_active = True + self.assertTrue(inner.is_active) + # Pre-trigger history is cleared on activation + for token in WS_IDS: + guard.feed(token) + guard.is_active = False + guard.is_active = True + self.assertFalse(guard._stalled) + + def test_reset_clears_state(self): + inner = FakeInner(packed_mask({1, 2, 6})) + guard = self.guard(inner) + for token in WS_IDS: + guard.feed(token) + guard.reset() + self.assertEqual(inner.reset_count, 1) + self.assertFalse(guard._stalled) + + def test_attach_reaches_inner(self): + inner = FakeInner(packed_mask({1, 2, 6})) + guard = self.guard(inner) + job = SimpleNamespace(generator=SimpleNamespace(padded_vocab_size=32)) + guard.attach(job) + self.assertIs(inner.attached, job) + + def test_dense_mask_path(self): + inner = FakeInner(dense_mask({1, 2, 3, 6})) + guard = self.guard(inner) + for token in [*WS_IDS, *WS_IDS]: # 8 ws tokens + guard.feed(token) + mask = guard.get_next_logit_mask() + allowed = {i for i in range(VOCAB) if torch.isfinite(mask[0, i])} + self.assertEqual(allowed, {1, 6}) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file