Skip to content

design.md

Latest commit

 

History

History
254 lines (187 loc) · 8.17 KB

design.md

File metadata and controls

254 lines (187 loc) · 8.17 KB
Raw

Design Document — guided-decoding-bench

Overview

This project benchmarks the per-token latency overhead of constrained decoding strategies for structured output generation in LLM serving.

Three strategies are compared:

  1. Free decoding — no constraints, argmax from raw logits
  2. Token healing — free decoding followed by one-shot JSON repair
  3. FSM guided decoding — per-token logit masking via a finite-state machine

The central question is:

How much does enforcing structural constraints cost per generated token, and where exactly does that cost come from?

Motivation

Most production LLM serving systems enforce structured outputs such as JSON, function call arguments, or schema-validated fields. The mechanism is typically logit masking: at each decode step, tokens that would violate the current grammar state are zeroed out before sampling.

The cost of this masking is not obvious. It depends on:

  • the size of the vocabulary
  • the number of valid tokens at each state
  • the complexity of the grammar representation
  • whether state transitions are cheap or expensive

This project isolates each component and measures it independently.

Strategies benchmarked

Free decoding

The simplest baseline. No constraints are applied. Tokens are chosen by argmax over raw logits. The output may or may not be valid JSON.

Token healing

Free decoding is performed without constraints. After generation completes, a one-shot post-processing step extracts the longest valid JSON prefix from the output. No per-token overhead during decode. Does not guarantee online structural validity.

FSM guided decoding

A finite-state machine is compiled from the target schema. At each decode step:

  1. The FSM provides the set of valid next characters
  2. A logit mask is built from this set
  3. The mask is applied to logits before argmax
  4. The FSM state is advanced with the chosen character

This enforces structural constraints at every step, online.

Architecture

src/config.py

Defines sweep parameters: schema types, string lengths, confidence levels, number of trials.

src/vocab.py

Character-level vocabulary: JSON structural characters, lowercase letters, digits, and space. Provides encode and decode functions.

src/schemas.py

Generates random valid JSON examples for three schema families:

  • flat: single-level object with string, int, and bool fields
  • nested: object containing a nested object
  • array: object containing an integer array

src/oracle_model.py

Synthetic logits generator. Given a gold output string and the current position, generates logits where the correct next character receives a boosted score and all others receive Gaussian noise. Confidence parameter controls the signal-to-noise ratio. This isolates guided decoding overhead from actual model inference cost.

src/fsm.py

FSM construction and state management.

Each schema is compiled to a sequence of segments:

  • LITERAL: fixed character sequence
  • WORD: variable-length lowercase alphabetic string
  • INT: fixed-width integer
  • BOOL: literal true or false

Each FSMState caches its valid-token mask as a pre-built tensor. The warmup_masks method pre-materializes all masks on the target device before the timed decode loop begins. This separates mask cache population from the per-step masking cost measured during benchmarking.

src/free_decode.py

Implements free decode loop with CUDA event timing. Computes valid JSON rate via json.loads rather than exact match.

src/healing.py

Implements free decode followed by longest-valid-JSON-prefix extraction. Measures heal time separately from decode time.

src/guided_decode.py

Implements FSM-guided decode loop with per-component timing:

  • mask_build_ms: time to retrieve the prebuilt mask for the current state
  • model_ms: time to generate logits
  • mask_apply_ms: time to apply mask and select token
  • state_update_ms: time to advance FSM state

src/benchmark.py

Orchestrates one configuration: builds FSM, runs warmup, executes trials, aggregates results. Measures both cold-build and warm-build FSM construction time separately.

src/metrics.py

Computes summary statistics including slowdown ratios, component percentages, and validity rates.

src/analysis.py

Generates plots: slowdown by schema, slowdown by confidence, component breakdown, valid JSON rates, slowdown by string length. Writes summary file.

Key design decisions

Character-level vocabulary

Using a character-level vocabulary rather than a real tokenizer simplifies the FSM construction significantly and removes tokenizer overhead from the measurement. This makes the benchmark focus cleanly on masking cost rather than tokenization complexity.

Oracle synthetic model

Using synthetic logits rather than a real transformer model means the model-forward cost is controlled and repeatable. This lets the benchmark isolate guided-decoding overhead precisely. In a production system, the model forward would dominate and the guided-decoding overhead would be a smaller fraction.

Prebuilt mask cache

FSMState caches its mask tensor per device. All masks are materialized before the timed decode loop via warmup_masks. This ensures that the measured mask_build_ms reflects only the time to retrieve the cached mask, not the time to construct it from scratch. Cold construction is measured separately as fsm_build_cold_ms.

Cold vs warm build separation

The first FSM build on a new device includes CUDA allocator and tensor runtime warmup overhead. Subsequent builds are much faster. Reporting both separately gives a clearer picture of steady-state serving cost.

Validity metric

Valid output is defined as json.loads succeeding, not as exact match with the gold string. This is the operationally correct definition for serving use cases.

Sweep parameters

Schema types: flat, nested, array String lengths: 4, 8, 16 Confidence levels: 0.6, 0.8, 0.95 Trials per config: 30

Total configurations: 27

Results summary

Mean latency per decode step: free: 0.234 ms healing: 0.232 ms fsm: 0.459 ms

Slowdown vs free baseline: healing: 0.993x mean fsm: 1.958x mean 1.948x median 2.109x max

FSM step breakdown: model logits: 14.9% mask build: 36.8% mask apply: 47.4% state update: 0.2%

One-time FSM build: cold build mean: 33.0 ms warm build mean: 22.0 ms

Valid JSON output rate: free: 0.933 healing: 0.929 fsm: 0.784

Key findings

  1. Token healing is effectively free during decode. The post-processing step adds negligible cost during the generation loop. The tradeoff is that it cannot enforce structural constraints online.

  2. FSM guided decoding adds roughly 2x per-token overhead. This is consistent across schema types, string lengths, and confidence levels.

  3. The dominant cost is mask construction and mask application. Grammar state updates account for only 0.2% of FSM step time. Mask build and mask apply together account for over 84% of FSM overhead.

  4. The overhead is deterministic with respect to model confidence. Changing oracle confidence from 0.6 to 0.95 does not change the FSM slowdown meaningfully. The masking cost is structural, not model-dependent.

  5. Cold-start FSM build time includes runtime warmup effects. The first build is 33ms mean but warm builds are 22ms mean. For long-running serving workers, the relevant cost is the warm build.

  6. Validity improvements are schema-dependent. Flat and nested schemas show validity rates close to free decoding. Array schemas show lower validity rates due to stricter integer sequence constraints in the oracle setup.

Scope and limitations

This benchmark measures the overhead of constrained decoding in a controlled synthetic setting. The following are not modeled:

  • Real transformer forward pass cost
  • BPE tokenizer integration
  • Multi-token constraints and prefix trees
  • Batch masking across concurrent requests
  • Full CFG or Earley parser strategies
  • Production caching of grammar states across requests

The results should be interpreted as an isolation benchmark for masking overhead, not as a complete production serving benchmark.

Author

Joao Felipe De Souza 2026