🛡️ Fail closed on un-compilable grammar constraints; re-rail whitespace stalls

backends/exllamav3/grammar.py

@@ -6,16 +6,61 @@
     Filter,
     LLGuidanceFilter,
 )
+from backends.exllamav3.whitespace_guard import (
+    WhitespaceStallGuardFilter,
+    build_whitespace_bitmask,
+)
+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."""
 
     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,
@@ -36,16 +81,17 @@ 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)
+        self.filters.append(self._wrap_with_stall_guard(lmfilter, tokenizer))
 
     def add_regex_filter(
         self,
@@ -62,16 +108,17 @@ 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)
+        self.filters.append(self._wrap_with_stall_guard(lmfilter, tokenizer))
 
     def add_grammar_filter(
         self,
@@ -94,13 +141,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)
+        self.filters.append(self._wrap_with_stall_guard(lmfilter, tokenizer))

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:

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

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."""
 

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

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? "