Fast JSON decoding of native LLM using compressed finite state machines

Fast JSON decoding of native LLM using compressed finite state machines

Author: Liangsheng Yin, Ying Sheng, Lianmin Zheng
Date: February 5, 2024

The content of this article is based on a blog article published by LMSYS Org. Original link:LMSYS Org Blog. The relevant code base can be found at the following link:SGLang code base。

Having an LLM always produce valid JSON or YAML that conforms to a specific schema is a critical feature for many applications. In this blog post, we introduce an optimization method that significantly speeds up the decoding of such constraints. Our approach leverages compressed finite state machines and is compatible with any regular expression and therefore can be applied to any JSON or YAML schema. Unlike existing systems that decode one token step by step, our approach analyzes a finite state machine of a regular expression, condenses a single transformation path, and decodes multiple tokens at once when possible. Compared to state-of-the-art systems (guidance + llama.cpp, outlines + vLLM), our approach can reduce latency by up to 2x and improve throughput by up to 2.5x. This optimization also makes constrained decoding faster than normal decoding. You can try it out on SGLang.

Figure 1 shows the performance comparison of SGLang and Outlines + vLLM in the JSON decoding task. This is a dynamic comparison, the purpose is to show the speed difference between the two under the same task. SGLang uses a new skip-forward decoding algorithm to speed up the decoding process by compressing the finite state machine. In contrast, Outlines + vLLM uses a traditional stepwise decoding approach. The animation in the figure demonstrates the advantages of SGLang when handling multi-character (or token) decoding, significantly reducing decoding time.

background

JSON is one of the most important formats for data exchange. Requiring LLM to always generate valid JSON allows LLM's output to be easily parsed in a structured manner. Recognizing its importance, OpenAI introduced JSON Schema, which constrains models to always return valid JSON objects. However, more fine-grained control is often required to ensure that the generated JSON objects conform to a specific schema, for example:

Figure 2 shows an example of constrained generation, using large language models (LLMs) to generate objects that conform to a specific JSON schema. In this example, the JSON schema on the left defines an object, which contains three properties: name, age, and house, which are string and integer types respectively. The right side shows the output object of restricted generation. The model uses constraint generation technology to generate specific instances that match these attributes, such as the name of "Harry", the age of 15 years old, and the house belonging to "Gryffindor". This demonstrates the capabilities of LLMs in generating structured data while ensuring that the generated content conforms to a predetermined format.

For local LLM, there are two main ways to guide the model to generate JSON objects that conform to a specific schema.

Method 1: Based on finite state machines

This approach involves converting a JSON pattern into a regular expression. We can then build a finite state machine (FSM) based on regular expressions. FSM is used to guide the generation of LLM. In each state of the FSM, we can calculate the allowed transitions and identify the acceptable next token. This allows us to keep track of the current state during decoding and filter out invalid tokens by applying a logit bias to the output. You can learn more about this approach in the outlines paper.

Figure 3 shows how to use a finite state machine (FSM) to implement restricted decoding. In this process, the JSON schema is first converted into regular expressions, and then FSM is used to guide the generation of LLM. In the figure, the FSM state diagram shows the restricted generation process of the age field, where only legal numbers (such as 0-9) are allowed. The transitions for each state are defined by regular expression rules, ensuring that the generated JSON data is always valid. This approach controls the LLM to generate specific outputs by imposing constraints on the generation process.

The FSM approach leverages generalized regular expressions to define low-level rules that can be applied to a wide range of syntaxes, such as JSON schemas, IP addresses, and emails.

Limitations:

Since the FSM is built at the tag level, it can only transition states through one tag at each step. Therefore, it can only decode one token at a time, resulting in slower decoding.

Method 2: Based on interleaving

An alternative to converting the entire JSON pattern to a regular expression is to use interleave-based decoding. In this approach, a given JSON schema can be decomposed into several parts, each part containing a chunked pre-population part or a constrained decoding part. These different parts are interleaved and executed by the inference system. Since chunked prefill can handle multiple tags in one forward pass, it is faster than tag-by-tag decoding.

Guidance provides a set of syntax rules based on interleaved decoding, using llama.cpp as the backend.

Figure 4 shows the interleaved syntax in the Guidance framework and how to use the interleaved syntax to decode JSON. The code snippet in the figure defines a function that uses Guidance syntax to generate a JSON object containing name, age, and house. Interleaved syntax can increase decoding speed by alternating decoding of different parts with pre-padding parts. The lower part of the figure shows how this process works. The green and blue bars represent different parts of the processing, showing the execution of interleaved decoding at different stages.

Limitations:

• The interleaving-based approach requires a custom syntax, making it less flexible and expressive than a single regular expression.
• There are difficulties in handling marker boundaries due to possible conflicts between decoded and chunked prefill segments.
• Frequent communication between the interpreter and the backend brings additional overhead.

Our approach: skip-forward decoding using compressed finite state machines

By introducing a new decoding algorithm based on compressed finite state machines - skip-forward decoding, we can combine the advantages of FSM and interleaving methods.

During the decoding process guided by the regular expression of the JSON pattern transformation, it is possible to predict the upcoming string when we reach a specific node:

In Figure 3, when decoding begins, based on the regular expression, we can predict that the next string is:

1
2

{
  "name":

Then onto the actual decoding part.
Likewise, when LLM outputs G when filling in the house properties for the character, we can confidently predict that the next string will be ryffindor, thus completing the entire string as Gryffindor.

This is exactly how the skip-forward decoding algorithm speeds up decoding. In the skip-forward algorithm, we examine a finite state machine given a regular expression, identify all single transition edges, and condense consecutive transition paths into a single path. We can directly prepopulate (extend) these single paths, skipping token-by-tag decoding until the next branch point.

Figure 5 shows the comparison between skip forward decoding and ordinary decoding. Skip-forward decoding utilizes compressed finite state machines to reduce the number of token-by-tag decodes by predicting and pre-populating possible strings in advance. For example, when generating a value for the house field, the model jumps directly into the decoding process and pre-populates the string "Gryffindor" without generating it character by character. The flow in the figure shows how to improve decoding efficiency through this method while avoiding unnecessary repeated calculations.
Figure 5 showsJump-forward decoding of compressed finite state machineswithNormal decodingComparison, especially the performance difference when generating JSON data. In order to understand this diagram in more detail, we need to analyze each part of the diagram step by step.

1. Input prompt(green part on the left): Prompts the model to generate an object that conforms to the JSON schema. The JSON object here includes three attributes: "name", "age" and "house", which represent name, age and college respectively.

2. Jump forward decoding process(blue and orange squares in the middle):

   • orange squareRepresents the part that requires constraint decoding. For example, when generating the "name" attribute, the model can directly generate the complete string "Harry" through the skip-forward decoding algorithm.
   • blue squareRepresents the portion of the model that decodes character by character (or token by token) during the jump forward process. This type of decoding only occurs when non-determinism is encountered (such as multiple possible values).

3. Ordinary decoding process(blue square in the middle): Normal decoding requires generating the entire JSON object character by character or token by token. In contrast, ordinary decoding methods require prediction and selection when processing each character or mark, which significantly reduces the decoding speed.

4. Compare results(right part):

   • Jump forward decodingThe generated JSON object is shown at the top. This approach greatly speeds up the decoding process by predicting and pre-populating possible strings. For example, when generating the string "Gryffindor", the model skips the character-by-character generation step.
   • Normal decodingThe generated JSON object is shown at the bottom. This method decodes character by character. Although it can ensure the accuracy of the generation, it is less efficient, especially when dealing with long strings or complex structures.

Detailed interpretation:

1. How skip-forward decoding works：

   • During the decoding process, the model uses a compressed finite state machine (FSM) to predict and identify the upcoming string. If the model can accurately predict the strings it will generate next in the current context, it can skip token-by-tag decoding of those strings and directly generate the entire string (e.g. "Gryffindor").
   • This method takes advantage of the structural characteristics of regular expressions to compress continuous transformation paths into a single path, thereby avoiding unnecessary token-by-tag decoding steps.

2. Limitations of ordinary decoding：

   • The normal decoding method requires decoding each character or token step by step and is therefore less efficient when dealing with complex JSON objects. Each step requires the model to recalculate possible outputs and select the optimal solution among them, which significantly increases decoding time.

3. Performance differences：

   • Since skip-forward decoding reduces the number of character-by-character decodings and takes advantage of the compression properties of FSM, it is significantly less expensive in terms of time and computing resources than ordinary decoding. Especially when large amounts of data need to be generated or complex structures need to be processed, the advantages of jump-forward decoding are even more obvious.

SGLang's RadixAttention mechanism greatly simplifies the implementation of skip-forward decoding algorithms. When executing a jump forward, we can simply terminate the current request and enqueue a new request. The SGLang runtime's RadixAttention and efficient extension primitives will automatically reuse the KV cache from the previous set of tags, thus avoiding redundant computations.

Mark boundary processing

When implementing constrained decoding, handling token boundaries is always tricky due to the complex possible mappings between characters and tokens.

During LLM decoding, it may be preferable (meaning higher probability) to combine multiple characters into a single token. For example, when decoding "Hello" in the context of JSON decoding, LLM might output the following token:
“ He llo “，

Instead of decoding the final " , it always tends to combine it with subsequent characters into the more common token " , this effect can lead to some strange behavior. For example, in the above case, if the regex is set to "[\w\d\s]*" (excluding the last ", ), this may result in infinite decoding because LLM wants to stop at ", but that mark is not allowed.

Furthermore, during skip-forward decoding, we found that using different marking strategies for the skip-forward part may result in different logit distributions for subsequent markings. Simply appending the tokenized jump-forward portion to the current token sequence may have unexpected results.

To address these issues, we propose the following solutions:

• We implemented a re-tokenization mechanism during the leap forward phase. This involves appending strings instead of tokens, and then re-tokenizing the entire text. This approach effectively solves most tokenization problems and only results in an increase of about 4% in computational overhead.
• SuggestionsUse a comprehensive regular expression to guide the entire decoding process instead of using multiple concatenated regular expressions. This approach ensures that both FSM and LLM understand the entire decoding process, thus minimizing boundary-related issues.
  You can also read some additional discussion in this blog post.

Benchmark results

We benchmarked our skip-forward decoding on two tasks:

1. Generate role data in JSON format using a short prompt.
2. Extract city information from long documents and output in JSON format.

We tested llama-7B on an NVIDIA A10 GPU (24GB), using vllm v0.2.7, guidance v0.1.0, outlines v0.2.5 and llama.cpp v0.2.38 (Python bindings). The following graph shows the throughput (using the maximum batch size supported by each system) and latency (batch size 1) for these methods:

The results show that SGLang using our decoding algorithm significantly outperforms all other systems. It can reduce latency by up to 2x and increase throughput by up to 2.5x. In the character generation task, SGLang achieved higher throughput than Outlines+vLLM even without using jump-forward; we suspect this is due to some overhead in Outlines.

use case

We have been testing this feature with Boson.ai for two weeks and they are in the process of introducing this feature into their production use cases as it guarantees higher decoding throughput and reliable response.

Additionally, another user used this feature to extract structured information from images via the visual language model LLaVA.