Repository: vllm-project/vllm

Issue:
Support for grammar
It would be highly beneficial if the library could incorporate support for Grammar and GBNF files.
https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md


Discussion:
Support this request. Would be massive to produce structured data at scale.
I will test whether [Shopify/torch-grammar](https://github.com/Shopify/torch-grammar) works around [these lines](https://github.com/vllm-project/vllm/blob/f8a1e39fae05ca610be8d5a78be9d40f5274e5fc/vllm/model_executor/layers/sampler.py#L83C14-L83C14).

Will be posting here if it works :)
Any updates?
Hi! Sorry for the delay. I ended up implementing a minor speedup for constrained generation on llama.cpp by skipping "obvious" tokens and just decoding the ones that _need_ an LLM. Although it managed to reach around double the original decoding performance, it is still not enough for my use case, so I will likely revisit this approach with vLLM in the future. 
It is not the current priority for my personal project though, so it may take some weeks before I start experimenting again. 

If anyone is eager to try an implementation, the idea I had in mind is, on a high level, as follows:
1. Instantiate the grammars, one for each sequence that will need to be decoded in parallel. I think it should be an attribute of the [Sampler](https://github.com/vllm-project/vllm/blob/f8a1e39fae05ca610be8d5a78be9d40f5274e5fc/vllm/model_executor/layers/sampler.py#L17) class:

    ```python
    class Sampler(nn.Module):


        def __init__(self, vocab_size: int) -> None:
            super().__init__()
            self.vocab_size = vocab_size
            # Dict[int, GrammarSampler]. Maps a sequence_id to a fresh sampler
            self.grammars = defaultdict(lambda: GrammarSampler(input_text, "root", tokenizer))


        def forward(...):

    ``` 

    This follows the example from `torch-grammar`'s [Readme](https://github.com/Shopify/torch-grammar/blob/main/README.md). As I'm writing this, I don't see how to know how many grammar samplers to create in advance, so the idea is to use a `defaultdict` that will create a new grammar as soon as an unseen `sequence_id` needs it.
2. Somehow, we should keep track of all the token IDs generated so far, because they are needed in the grammar sampler to set the logits of invalid continuation tokens to zero, and thus prevent those tokens from being randomly sampled and selected. A great candidate for this seems to be the `output_tokens = _get_output_tokens(input_metadata)` call a few lines into the `forward` function. This holds a list of tokens ids for each sequence id. It's enlightening to read the [definition](https://github.com/vllm-project/vllm/blob/f8a1e39fae05ca610be8d5a78be9d40f5274e5fc/vllm/model_executor/layers/sampler.py#L157) of `_get_output_tokens` towards the middle of the file to know the general shape of the returned list.
3. Just [before computing the `logprobs`](https://github.com/vllm-project/vllm/blob/06e9ebebd51c3db779dedec5556251c8ecc3a00a/vllm/model_executor/layers/sampler.py#L91), we could call the `logits_processor` for each pair of `(seq_id, grammar)`. The important thing here is to match up a few things: each sequence id to the corresponding grammar, that keeps track of the state of the generation and defines which tokens will be discarded now, and also to the current output for each of those sequence ids.


    ```python
    def forward(...):
        # More sampling code above ...

        # We use float32 for probabilities and log probabilities.
        # Compute the probabilities.
        probs = torch.softmax(logits, dim=-1, dtype=torch.float)
        
        # I don't know how to vectorize this, so I unroll it in a simple for loop
        output_idx = 0
        for seq_ids, sampling_params in input_metadata.seq_groups:
            for seq_id in seq_ids:
                #  the defaultdict will create a clean grammar when a given seq_id doesn't exist yet
                grammar_sampler = self.grammars[seq_id]
                output_for_this_seq_id = output_tokens[output_idx]
                logits_for_this_seq_id = logits[output_idx]
                # call the logits_processor(ids, logits)
                logits[output_idx] = grammar_sampler.logits_processor(output_for_this_seq_id,
                                                                      logits_for_this_seq_id)
                # ^ this will zero out the logits that correspond to invalid token ids under the grammar
                output_idx += 1
                # Note that I deliberately don't use an `enumerate` to update the output_idx variable,
                # because the _get_output_tokens function appends a sequence once for each seq_id,
                # not on every seq_group!
 
        # Continue with the existing code

        # Compute the log probabilities.
        # Use log_softmax to ensure numerical stability.
        logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)
    ``` 
4. Accept the tokens that have been sampled to update the grammar state for each sequence id. Each sampling strategy (greedy, random, beam) uses the same interface for the return: `List[Tuple[List[int], List[int]]]`. Refer to the implementation of `_greedy_sample`, `_random_sample` or `_beam_sample` for more details, but from what I can understand by skimming through the code, those functions return, for each sequence group that requested that specific sampling strategy (this is the outer `List` of the return type), a tuple of the form `(next_token_ids, parent_ids)` to support generating multiple completions (one new token id for each `sequence_id`) for any single sequence group. Apparently, this tracks which tokens are selected (in parallel) as the next id and which ones were responsible for that generation (`parent_ids`). 
I know this is like the movie Inception, but hold on. 

    ```python
    def forward(...):
        # More sampling code above ...

        # Compute the log probabilities.
        # Use log_softmax to ensure numerical stability.
        logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)
        
        # Sample the next tokens.
        sample_results = _sample(probs, logprobs, input_metadata)

        # Now we accept the token for each sequence id inside every sequence group
        output_idx = 0
        for seq_group_id, (seq_ids, sampling_params) in enumerate(input_metadata.seq_groups):
            next_token_ids, parent_ids = sample_results[seq_group_id]
            for seq_id in seq_ids:
                grammar_sampler = self.grammars[seq_id]
                next_token_for_this_seq_id = next_token_ids[output_idx]
                grammar_sampler.accept_token(next_token_for_this_seq_id)
                output_idx += 1
 
        # Continue with the existing code 
    ``` 

As a final note, I'm not well familiarized with this project, and this plan is only a high level review of how I would go about implementing grammars. Also, note that the `torch-grammar` library is clearly unmaintained, but I think it should work well enough as a proof of concept. At least its API is comfortable.

Another update. I wasn't aware of the existence of `logits_processors`, as implemented on [this PR](https://github.com/vllm-project/vllm/pull/1469). That direction seems way more promising and already has one [example in the form of a unit test](https://github.com/noamgat/vllm/blob/417fc0033885602fcbbe19046b921585f90b33a6/tests/samplers/test_sampler.py#L189).
@viantirreau 

Take a look at https://github.com/noamgat/lm-format-enforcer

They're using logits_processors to enforce JSON schema's and regex = basically Grammar
@viantirreau There is a LM Format Enforcer Integration with vLLM colab tutorial on https://github.com/noamgat/lm-format-enforcer/blob/main/samples/colab_vllm_integration.ipynb 

Perhaps this can be use as guideline for future integrations
What is the technical reason why we cannot port over [GBNF](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md)? It could be a `logits_processor`, plus a grammar need to be passed via API requests (both vLLM and OpenAI).

Resolving patch:
diff --git a/vllm/entrypoints/openai/protocol.py b/vllm/entrypoints/openai/protocol.py
--- a/vllm/entrypoints/openai/protocol.py
+++ b/vllm/entrypoints/openai/protocol.py
@@ -55,6 +55,11 @@ class UsageInfo(BaseModel):
     completion_tokens: Optional[int] = 0
 
 
+class ResponseFormat(BaseModel):
+    # type must be "json_object" or "text"
+    type: str = Literal["text", "json_object"]
+
+
 class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[Dict[str, str]]
@@ -89,6 +94,8 @@ class ChatCompletionRequest(BaseModel):
     guided_json: Optional[Union[str, dict, BaseModel]] = None
     guided_regex: Optional[str] = None
     guided_choice: Optional[List[str]] = None
+    guided_grammar: Optional[str] = None
+    response_format: Optional[ResponseFormat] = None
 
     def to_sampling_params(self) -> SamplingParams:
         if self.logprobs and not self.top_logprobs:
@@ -183,6 +190,8 @@ class CompletionRequest(BaseModel):
     guided_json: Optional[Union[str, dict, BaseModel]] = None
     guided_regex: Optional[str] = None
     guided_choice: Optional[List[str]] = None
+    guided_grammar: Optional[str] = None
+    response_format: Optional[ResponseFormat] = None
 
     def to_sampling_params(self):
         echo_without_generation = self.echo and self.max_tokens == 0
diff --git a/vllm/model_executor/guided_decoding.py b/vllm/model_executor/guided_decoding.py
--- a/vllm/model_executor/guided_decoding.py
+++ b/vllm/model_executor/guided_decoding.py
@@ -6,19 +6,50 @@
 from json import dumps as json_dumps
 from re import escape as regex_escape
 from typing import Union, Tuple
+
 from pydantic import BaseModel
+from transformers import PreTrainedTokenizerBase
 
 from vllm.entrypoints.openai.protocol import (CompletionRequest,
                                               ChatCompletionRequest)
 from vllm.model_executor.guided_logits_processors import (JSONLogitsProcessor,
-                                                          RegexLogitsProcessor)
+                                                          RegexLogitsProcessor,
+                                                          CFGLogitsProcessor)
 
 
 class GuidedDecodingMode(Enum):
     JSON = "json"
     REGEX = "regex"
     CHOICE = "choice"
+    GRAMMAR = "grammar"
+
+
+# https://github.com/outlines-dev/outlines/blob/main/outlines/grammars/json.lark
+# the main difference is that we changed the start: value to
+# start: object | array, so we are denying scalar values as the root of the
+# JSON. Starting with scalars as the root seems to cause llama to generate
+# without stop.
+JSON_GRAMMAR = r"""
+?start: object | array
+
+?value: object
+| array
+| UNESCAPED_STRING
+| SIGNED_NUMBER      -> number
+| "true"             -> true
+| "false"            -> false
+| "null"             -> null
+
+array  : "[" [value ("," value)*] "]"
+object : "{" [pair ("," pair)*] "}"
+pair   : UNESCAPED_STRING ":" value
+
+%import common.UNESCAPED_STRING
+%import common.SIGNED_NUMBER
+%import common.WS
 
+%ignore WS
+"""
 
 global_thread_pool = None  # used for generating logits processor fsm
 
@@ -57,9 +88,6 @@ def _get_guide_and_mode(
 ) -> Tuple[str, GuidedDecodingMode]:
 
     if request.guided_json:
-        if not isinstance(request.guided_json, (str, dict, BaseModel)):
-            raise TypeError("JSON schema must be str, dict, or BaseModel")
-
         json = request.guided_json
         if isinstance(json, dict):
             # turn dict into hashable string
@@ -69,33 +97,33 @@ def _get_guide_and_mode(
             # with the same fields will get hashed the same
             json = str(json.__signature__)
         return json, GuidedDecodingMode.JSON
-
     elif request.guided_regex:
-        if not isinstance(request.guided_regex, str):
-            raise TypeError("Regex must be string")
         return request.guided_regex, GuidedDecodingMode.REGEX
-
     elif request.guided_choice:
-        if not isinstance(request.guided_choice, list):
-            raise TypeError("Choices must be a list")
-
         # choice just uses regex
         choices = [
             regex_escape(str(choice)) for choice in request.guided_choice
         ]
         choices_regex = "(" + "|".join(choices) + ")"
         return choices_regex, GuidedDecodingMode.CHOICE
-
+    elif request.guided_grammar:
+        return request.guided_grammar, GuidedDecodingMode.GRAMMAR
+    elif (request.response_format is not None
+          and request.response_format.type == "json_object"):
+        return JSON_GRAMMAR, GuidedDecodingMode.GRAMMAR
     else:
         return None, None
 
 
 @lru_cache(maxsize=32)
-def _get_cached_logits_processor(guide: str, tokenizer,
+def _get_cached_logits_processor(guide: str,
+                                 tokenizer: PreTrainedTokenizerBase,
                                  mode: GuidedDecodingMode):
     if mode == GuidedDecodingMode.JSON:
         return JSONLogitsProcessor(guide, tokenizer)
     elif mode == GuidedDecodingMode.REGEX or mode == GuidedDecodingMode.CHOICE:
         return RegexLogitsProcessor(guide, tokenizer)
+    elif mode == GuidedDecodingMode.GRAMMAR:
+        return CFGLogitsProcessor(guide, tokenizer)
     else:
         raise ValueError(f"Unknown guided decoding mode {mode}")
diff --git a/vllm/model_executor/guided_logits_processors.py b/vllm/model_executor/guided_logits_processors.py
--- a/vllm/model_executor/guided_logits_processors.py
+++ b/vllm/model_executor/guided_logits_processors.py
@@ -16,30 +16,60 @@
 import json
 import math
 from collections import defaultdict
-from typing import Union, DefaultDict, Dict, List, Optional
+from typing import Union, DefaultDict, Dict, List, Optional, Callable
 
 import torch
 from pydantic import BaseModel
-from outlines.fsm.fsm import RegexFSM
+from transformers import PreTrainedTokenizerBase
+from outlines.fsm.fsm import RegexFSM, CFGFSM
 from outlines.fsm.json_schema import build_regex_from_schema
 
 
-class RegexLogitsProcessor:
+class BaseLogitsProcessor:
 
-    def __init__(self, regex_string: str, tokenizer):
-        """Compile the FSM that drives the regex-structured generation.
+    def adapt_tokenizer(self, tokenizer: PreTrainedTokenizerBase):
+        """Adapt vLLM's tokenizer to use to compile the FSM.
 
-        Parameters
-        ----------
-        regex_string
-            A string that represents a regular expression
-        tokenizer
-            The model's tokenizer
+        The API of Outlines tokenizers is slightly different to that of
+        `transformers`. The decoder of outlines, returns a list whereas
+        the decode of vLLM returns an str. To sync the vLLM decoder with
+        outlines internal api, the decoder should be adapted. In addition
+        we need to handle the missing spaces to Llama's tokenizer to be
+        able to compile FSMs for this model.
 
         """
-        tokenizer = self.adapt_tokenizer(tokenizer)
-        fsm = RegexFSM(regex_string, tokenizer)
-        self.fsm = fsm
+        if getattr(tokenizer, "_outlines_adapted", False):
+            return tokenizer
+
+        tokenizer.vocabulary = tokenizer.get_vocab()
+        tokenizer.special_tokens = set(tokenizer.all_special_tokens)
+
+        def convert_token_to_string(token: str) -> str:
+            from transformers.file_utils import SPIECE_UNDERLINE
+
+            string = tokenizer.convert_tokens_to_string([token])
+
+            # A hack to handle missing spaces to HF's Llama tokenizers
+            if token.startswith(SPIECE_UNDERLINE) or token == "<0x20>":
+                return " " + string
+
+            return string
+
+        def change_decoder(
+            decoder: Callable[[List[int]], str]
+        ) -> Callable[[List[int]], List[str]]:
+            """Sync vLLM's decoder with the outlines by returning list."""
+
+            def new_decoder(inp_tokens: List[int]) -> List[str]:
+                return [decoder(inp_tokens)]
+
+            return new_decoder
+
+        tokenizer.convert_token_to_string = convert_token_to_string
+        tokenizer.decode = change_decoder(tokenizer.decode)
+        setattr(tokenizer, "_outlines_adapted", True)  # noqa: B010
+
+        return tokenizer
 
     def init_state(self):
         """Initialize the FSM states."""
@@ -69,38 +99,30 @@ def __call__(self, input_ids: List[int],
 
         return scores
 
-    def adapt_tokenizer(self, tokenizer):
-        """Adapt vLLM's tokenizer to use to compile the FSM.
-
-        The API of Outlines tokenizers is slightly different to that of
-        `transformers`. In addition we need to handle the missing spaces to
-        Llama's tokenizer to be able to compile FSMs for this model.
-
-        """
-        tokenizer.vocabulary = tokenizer.get_vocab()
-        tokenizer.special_tokens = set(tokenizer.all_special_tokens)
-
-        def convert_token_to_string(token: str) -> str:
-            from transformers.file_utils import SPIECE_UNDERLINE
 
-            string = tokenizer.convert_tokens_to_string([token])
+class RegexLogitsProcessor(BaseLogitsProcessor):
 
-            # A hack to handle missing spaces to HF's Llama tokenizers
-            if token.startswith(SPIECE_UNDERLINE) or token == "<0x20>":
-                return " " + string
-
-            return string
+    def __init__(self, regex_string: str, tokenizer: PreTrainedTokenizerBase):
+        """Compile the FSM that drives the regex-structured generation.
 
-        tokenizer.convert_token_to_string = convert_token_to_string
+        Parameters
+        ----------
+        regex_string
+            A string that represents a regular expression
+        tokenizer
+            The model's tokenizer
 
-        return tokenizer
+        """
+        tokenizer = self.adapt_tokenizer(tokenizer)
+        fsm = RegexFSM(regex_string, tokenizer)
+        self.fsm = fsm
 
 
 class JSONLogitsProcessor(RegexLogitsProcessor):
 
     def __init__(self,
                  schema: Union[str, Dict, BaseModel],
-                 tokenizer,
+                 tokenizer: PreTrainedTokenizerBase,
                  whitespace_pattern: Optional[str] = None):
         """Compile the FSM that drives the JSON-guided generation.
 
@@ -130,3 +152,21 @@ def __init__(self,
                 f"the JSON Schema specification")
         regex_string = build_regex_from_schema(schema_str, whitespace_pattern)
         super().__init__(regex_string, tokenizer)
+
+
+class CFGLogitsProcessor(BaseLogitsProcessor):
+
+    def __init__(self, cfg: str, tokenizer: PreTrainedTokenizerBase):
+        """Compile the FSM that drives the context free grammar generation.
+
+        Parameters
+        ----------
+        cfg
+            A string that represents a context-free grammar
+        tokenizer
+            The model's tokenizer
+
+        """
+        tokenizer = self.adapt_tokenizer(tokenizer)
+        fsm = CFGFSM(cfg, tokenizer)
+        self.fsm = fsm


Test patch:
diff --git a/tests/entrypoints/test_openai_server.py b/tests/entrypoints/test_openai_server.py
--- a/tests/entrypoints/test_openai_server.py
+++ b/tests/entrypoints/test_openai_server.py
@@ -660,5 +660,55 @@ async def test_guided_decoding_type_error(server, client: openai.AsyncOpenAI):
             extra_body=dict(guided_regex=TEST_REGEX, guided_json=TEST_SCHEMA))
 
 
+async def test_response_format_json_object(server, client: openai.AsyncOpenAI):
+    resp = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=[{
+            "role":
+            "user",
+            "content": ('what is 1+1? please respond with a JSON object, '
+                        'the format is {"result": 2}')
+        }],
+        response_format={"type": "json_object"})
+
+    content = resp.choices[0].message.content
+    loaded = json.loads(content)
+    assert loaded == {"result": 2}, loaded
+
+
+async def test_guided_grammar(server, client: openai.AsyncOpenAI):
+    simple_sql_grammar = """
+start: select_statement
+
+select_statement: "SELECT" column "from" table "where" condition
+
+column: "col_1" | "col_2"
+table: "table_1" | "table_2"
+condition: column "=" number
+
+number: "1" | "2"
+"""
+
+    completion = await client.completions.create(
+        model=MODEL_NAME,
+        prompt=("Generate a sql state that select col_1 from "
+                "table_1 where it is equals to 1"),
+        temperature=1.0,
+        max_tokens=500,
+        extra_body=dict(guided_grammar=simple_sql_grammar))
+
+    content = completion.choices[0].text
+
+    # use Lark to parse the output, and make sure it's a valid parse tree
+    from lark import Lark
+    parser = Lark(simple_sql_grammar)
+    parser.parse(content)
+
+    # remove spaces for comparison b/c we removed them in the grammar
+    ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "")
+
+    assert content.strip() == ground_truth
+
+
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/tests/kernels/test_prefix_prefill.py b/tests/kernels/test_prefix_prefill.py
--- a/tests/kernels/test_prefix_prefill.py
+++ b/tests/kernels/test_prefix_prefill.py
@@ -36,8 +36,8 @@ def test_contexted_kv_attention(
         torch.cuda.manual_seed(0)
     torch.set_default_device(device)
 
-    # Need this, otherwise when we capture the graph the process for GPU 1 would run on both
-    # GPU0 and GPU1 and things would hang
+    # Need this, otherwise when we capture the graph the process for GPU 1 would
+    # run on both GPU0 and GPU1 and things would hang
     #
     # see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
     torch.cuda.set_device(device)