Efficient Grammar-Constrained Decoding via Parser Stack Classification

This is the replication package of the paper "Efficient Grammar-Constrained Decoding via Parser Stack Classification".

Requirements

You can set up a conda environment using the provided environment.yml file:

conda env create -f environment.yml
conda activate fast-constraint

This will install all necessary dependencies to run the benchmarking code with PSC. To run other grammar-constrained decoding methods, you may need to install additional packages, as specified in later sections.

Benchmarking different methods

To benchmark different grammar-constrained decoding methods, navigate to the benchmark directory and run the following commands.

Measure the latency of mask computation

To calculate the latency of mask computation for method $method on language $lang with the tokenizer $tokenizer, and generate output file $output, run:

python -m src.only_mask_latency --data datasets/$lang.jsonl --lark grammars/lark/$lang.lark --gbnf grammars/gbnf/$lang.gbnf --kbnf grammars/kbnf/$lang.kbnf --model $tokenizer --engine $method --output $output --cache_dir $cache

where $cache is a directory to store the preprocessing result of PSC. We provide all the preprocessing results used in our experiments, where the preprocessing results of programming languages are stored at cache/, and the preprocessing results of JSON schemas are stored at cache/json-schemas-$tokenizer, where $tokenizer should be one of gemma, llama and qwen.

To calculate the latency of mask computation for method $method on schema-conformant JSON with the tokenizer $tokenizer, and generate output file $output, run:

python -m src.only_mask_latency --data datasets/json-schemas.jsonl --is_json --model $tokenizer --engine $method --output $output --cache_dir $cache

You have to use the name fast_constraint for PSC.

Note that you may need to install additional packages for certain methods:

• For xgrammar, run pip install --upgrade xgrammar.
• For llguidance, run pip install --upgrade llguidance.
• For formatron, run pip install --upgrade formatron.
• For greatgramma, follow the instructions at https://github.com/ebmoon/alignment to install GreatGramma.

To extract the latency from the output file $latency, run:

python -m src.calculate_latency --input $latency

Measure the end-to-end throughput

To calculate the end-to-end throughput for method $method on language $lang with the model $model and batch size $batch, and generate output file $output, run the following command:

python -m src.vllm_inference --data datasets/$lang.jsonl --lark grammars/lark/$lang.lark --gbnf grammars/gbnf/$lang.gbnf --kbnf grammars/kbnf/$lang.kbnf --model $tokenizer --engine $method --output $output --cache_dir $cache --reference $latency --max_batch_size $batch

Note that you should also provide the result of the output file $latency from the mask latency measurement step.

File structure

There are two subdirectories: fast_constraint and benchmark. The fast_constraint directory contains the implementation of the proposed PSC method, and the benchmark directory contains the code for benchmarking different grammar-constrained decoding methods.

fast_constraint

This directory contains the implementation of the PSC method.

• partial_lexer.py: The implementation of the lexical preprocessing step, borrowed from GreatGramma, as described in Section 3.1 of the paper.
• classify.py: The implementation of the parser stack classification.
• runner.py: The supportive code to run PSC as a grammar-constrained decoding method.
• utils.py: Utility functions used in the PSC implementation.
• __init__.py: Initializes the fast_constraint package.

benchmark

This directory contains the necessary code to benchmark different grammar-constrained decoding methods.

• datasets: Contains dataset files used in the experiments.
• grammars: Contains grammar files in various formats (including Lark, GBNF and KBNF).
• src: Contains the source code for benchmarking different methods.
  • vllm_inference.py: Code to run VLLM with grammar constraints.
  • only_mask_latency.py: Code to measure latency of mask computation.
  • calculate_latency.py: Code to calculate latency in the experiments.
  • ...: Other supportive scripts for benchmarking.

Trade-off between preprocessing time and memory requirement

There is a magic constant called increase_for_minimization in the construction algorithm (in classify.py). The final DFA $\mathcal{A}$, in this implementation, is constructed by parts, and minimized whenever its size increases by a certain threshold, and increase_for_minimization is that threshold.

In practice, this parameter controls the memory requirement and preprocessing time (when the preprocessing time is significantly long): when this threshold is large, the minimization happens infrequently, so the preprocessing requires more memory but less time; when this threshold is small, the DFA is often minimized, so the preprocessing requires less memory but longer time.

You should adjust this parameter by yourself based on the memory usage you can afford. You can also try to find a better algorithm to construct the DFA with less memory requirement and preprocessing time, which is one future work of this paper to further improve its efficiency and scalability.