Compressed FSM：受约束解码的"跳跃压缩"

谁该读这一篇？ 想给 LLM 强制输出 JSON / regex / 特定格式、又关心 TPOT 的工程师。 前置阅读： 01-overview/00-prerequisites.md（Decode 是什么）；最好对 FSM 有点印象（高中或大学课）。 耗时： 40 分钟 学完能： 1. 解释朴素受约束解码（vocab masking）为什么慢； 2. 画出 FSM 状态压缩与 jump-ahead 的两层优化； 3. 看到一个 JSON Schema 知道它会被编译成怎样的 FSM； 4. 选对 grammar 后端（xgrammar / outlines / llguidance / 内置压缩 FSM）； 5. 排查"约束开了但还是输出格式错"的情况。

1. 问题陈述：朴素受约束解码

业务："要求 LLM 输出严格 JSON，schema 是 {name: string, age: number}。"

朴素做法：

for step in decode_loop:
    logits = model(...)                      # 全 vocab 的 logits
    mask = compute_mask(current_state, schema)  # 哪些 token 合法？1 维 vocab 长度的 bool
    logits[~mask] = -inf
    token = sample(logits)
    advance_state(token, schema)

问题：

compute_mask 慢。对复杂 schema，每步要遍历 schema 算"在当前状态下哪些 token 是合法 next"。复杂嵌套 JSON 几毫秒。
逐 token 推进。"name": " 这串字符在 FSM 里只有一条路径（必走），但还是要逐字 sample —— 9 步 forward。
mask 张量很大。vocab 大小 = 152k（Qwen2.5）；每步算 152k 个 bool 不是免费的。

SGLang 的 压缩 FSM 同时干掉这三个问题。

2. 解决方案：两层压缩 + jump-ahead

2.1 状态压缩

观察：很多 schema 状态有"唯一确定的 token 序列"——例如 "name": 这一串。 FSM 里这部分是链式：

状态 0 ─"─▶ 状态 1 ─n─▶ 状态 2 ─a─▶ 状态 3 ─m─▶ 状态 4 ─e─▶ 状态 5 ─"─▶ 状态 6 ─:─▶ 状态 7

把这个链一次性折成一条"跳跃边"：

状态 0 ─"name":─▶ 状态 7

运行时一次直接吐 "name": 整段 token id，跳过 6 次 forward。

这就是 jump-ahead（跳跃解码 / jump-forward decoding）。源码：outlines_jump_forward.py。

2.2 Vocab mask 缓存

观察：相同 FSM 状态在请求间会重复出现（同一个 JSON schema 跑 1 万次，状态总数有限）。为每个状态预先算好 vocab mask 张量（vocab_size 大的 bool 张量），按 LRU 缓存复用，每步直接查表。

效果：单次 mask 计算 O(vocab_size) → O(1) 表查询。

3. Jump-ahead 的具体算法

源码：outlines_jump_forward.py:62。

伪代码：

def init_state_to_jump_forward(regex):
    fsm = build_fsm(regex)
    jump_table = {}
    for state in fsm.states:
        if has_unique_next(state):                  # 只有一种 transition
            path = follow_chain(state)              # 一直跟到分叉
            jump_table[state] = JumpEdge(
                symbol=path.symbols,                # 跳的字符串
                next_state=path.end_state,
            )
    return jump_table

decode loop 里：

def decode_step():
    if state in jump_table:
        edge = jump_table[state]
        token_ids = tokenizer.encode(edge.symbol)
        write_to_output(token_ids)                  # 一口气吐多 token
        state = edge.next_state
        skip_forward = True
    else:
        # 正常采样
        logits = model(...)
        mask = state_mask_cache[state]
        logits[~mask] = -inf
        token = sample(logits)
        state = advance(state, token)

注意一个微妙点：jump-ahead 必须按 tokenizer 重新 tokenize——FSM 是字符级的，但模型按 token 推进。"name": 在不同 tokenizer 里可能是 1 个 token 或 3 个 token，编码出来才知道。

4. 一个完整例子：JSON Schema

schema:

{"type": "object", "properties": {
    "name": {"type": "string"},
    "age": {"type": "integer"}
}}

编译出来的 FSM（简化）：

s0  ─{─▶ s1
s1  ─"name":─▶ s2     ← 一条 jump 边
s2  ─"─▶ s3
s3  ─{string char}─▶ s3   ← 字符串内容（自由采样）
s3  ─"─▶ s4
s4  ─,─▶ s5
s5  ─"age":─▶ s6      ← 一条 jump 边
s6  ─{digit}─▶ s7
s7  ─{digit}─▶ s7
s7  ─}─▶ s8 (accept)

decode 过程：

步	输入 state	动作	输出	新 state
1	s0	jump	`{`	s1
2	s1	jump	`"name":`	s2
3	s2	jump	`"`	s3
4	s3	sample（受限 string char）	`A`	s3
5	s3	sample	`l`	s3
...	...	...	...	...
...	s3	sample → 选了 `"`	`"`	s4
...	s4	jump	`,`	s5
...	s5	jump	`"age":`	s6
...	s6	sample（受限 digit）	`2`	s7
...	s7	sample	`8`	s7
...	s7	sample → 选了 `}`	`}`	s8

可以看到：

8 步里只有 ~4 步真正调了 model forward（A、l、...、2、8），其它都是 jump。
4 步 forward 完成了 30+ 个字符的输出。

TPOT 收益：30 字符 / 4 步 ≈ 7.5×。

5. 后端选择：4 种 grammar backend

SGLang 默认是自家压缩 FSM，但也集成了第三方：

后端	文件	优势	劣势
outlines + 压缩 FSM	`outlines_backend.py` + `outlines_jump_forward.py`	jump-ahead 强，简单 schema 极快	复杂 schema 编译慢
xgrammar	`xgrammar_backend.py`	编译快，mask 算得快	jump-ahead 较弱
llguidance	`llguidance_backend.py`	Microsoft 出品，支持复杂 grammar（CFG）	中文 / 多字节字符支持需注意
reasoner_grammar	`reasoner_grammar_backend.py`	reasoning 模型专用（think 标签等）	场景特化

选择策略：

简单 JSON：默认 outlines + jump-ahead 即可。
复杂深嵌 JSON / 长字符串字段：xgrammar 更稳。
需要 CFG（不只是 regex）：llguidance。
DeepSeek-R1 / QwQ 这类带 think 的模型：reasoner_grammar。

切换：

python -m sglang.launch_server --grammar-backend xgrammar

6. 客户端怎么用

Python (openai SDK)

resp = client.chat.completions.create(
    model="...",
    messages=[{"role": "user", "content": "返回 JSON"}],
    response_format={
        "type": "json_schema",
        "json_schema": {
            "name": "person",
            "schema": {...}
        }
    }
)

SGLang DSL

@sgl.function
def extract(s, text):
    s += sgl.user(text)
    s += sgl.assistant_begin()
    s += sgl.gen("person", json_schema=json.dumps({
        "type": "object",
        "properties": {
            "name": {"type": "string"},
            "age": {"type": "integer"}
        }
    }))
    s += sgl.assistant_end()

或者直接用 regex：

s += sgl.gen("color", regex=r"red|blue|green")

7. 编译时间不是免费的

复杂 schema 第一次见到时 SGLang 要编译 FSM，可能花 100ms-数秒。 SGLang 把编译结果缓存在 GrammarManager 里（grammar_manager.py），同 schema 再次出现直接复用。

工程上要注意：

同 schema 应该尽量使用相同字符串表示（不要随机字段顺序），否则缓存 miss。
业务上线前 warmup 一次，把热 schema 编译好。

8. 病态情况

情况	现象	解决
schema 太深嵌套	编译卡死 / OOM	简化 schema、用 `xgrammar`
字符串字段允许任意 Unicode	jump-ahead 退化	加 length / pattern 限制
Schema 里包含中文 literal	jump-ahead 可能错 byte	测一遍；考虑改用 llguidance
模型生成时跑出 schema 外	mask 强制约束应该不会 happen	极少数 bug，提 issue
受约束输出比无约束还慢	编译 / mask 没缓存	检查 schema 唯一性、第一次发预热

9. 性能数字（粗略）

schema	无约束 TPOT	xgrammar TPOT	SGLang 压缩 FSM TPOT
简单 JSON（2 字段）	12 ms	14 ms	9 ms（jump 多）
中等 JSON（10 字段嵌套）	12 ms	16 ms	13 ms
复杂 CFG	12 ms	20 ms	25 ms

SGLang 压缩 FSM 在"高 jump"场景下能比无约束更快（因为 jump 把多个字符一步吐完）；在复杂 CFG 上劣于 llguidance。

10. 调试 cheatsheet

# 加 grammar 调试日志
SGLANG_GRAMMAR_LOG_LEVEL=debug python -m sglang.launch_server ...

看 log 里：

"Grammar compile: X ms" 是不是过长。
"Jump forward: N tokens skipped" 看是不是 jump 真的起作用。
"Vocab mask cache hit/miss" 看缓存有没有命中。

11. 小结

压缩 FSM = 状态压缩（链折成 jump 边）+ vocab mask LRU 缓存。
Jump-ahead 让"唯一确定路径"一步吐多 token，省 forward。
选 backend：简单 JSON outlines、复杂 schema xgrammar、CFG llguidance。
编译有开销，业务上线前 warmup 一次。
高 jump 场景受约束输出甚至比无约束快。

12. 自检

朴素受约束解码的三个慢点分别是什么？压缩 FSM 怎么解决？

答案

慢点：(a) 每步 `compute_mask` 算"哪些 token 合法"——遍历 schema + tokenizer 反查，复杂 schema 几 ms；(b) 逐 token 推进 FSM，对"唯一确定 token 路径"也要每步 sample；(c) vocab mask 张量 = vocab_size 大的 bool，每步重算开销大。压缩 FSM 解决：(a) **预编译 + LRU 缓存 vocab mask 张量**（每状态一份，查表 O(1)）；(b) **Jump-ahead** 把链式状态折叠成跳跃边，一次吐多 token 跳过 sample；(c) cache 复用让重复 schema 不再编译。

Jump-ahead 在中文 schema 上可能失效，为什么？

答案

Jump-ahead 在 FSM 里是字符级的，但模型按 token 推进，跳的字符串必须**整段对齐 token 边界**。中文里"你好"可能是 1 个 token、也可能是"你"+"好"两个 token；不同模型 tokenizer 切法不同。Jump-ahead 计算时如果按字节直接 tokenize，可能切出"半个字符"或"半 token"，导致下一步 model forward 输入畸形。实现里有 byte-level 处理路径（[`outlines_jump_forward.py`](../sglang/python/sglang/srt/constrained/outlines_jump_forward.py)），但中文 / emoji / 复杂 Unicode 仍是 edge case 高发区，必要时 fallback 到逐 token 采样。

同 schema 多请求时编译为什么不会重复发生？

答案

`GrammarManager`（[`srt/constrained/grammar_manager.py`](../sglang/python/sglang/srt/constrained/grammar_manager.py)）维护一个 `schema_str -> compiled_grammar` 的字典缓存。关键前提：客户端发的 schema 字符串要**完全一致**（含字段顺序、空格）。`json.dumps(schema, sort_keys=True)` 是必备技巧；如果业务每次随机 dict 顺序，hash 不同会 cache miss。生产建议：业务侧预先编译好 schema 字符串作为常量，不要每请求重新 `json.dumps`。

xgrammar 比 outlines 快的常见场景是什么？

答案

xgrammar 在以下场景比 outlines + jump-ahead 快： (a) **复杂深嵌套 schema**：xgrammar 编译速度快得多（GPU 加速 mask 计算），outlines 编译 oneOf × N 分支会爆。 (b) **mask 计算 dominant**：自由字段多、jump-ahead 机会少时，xgrammar 的 GPU bitmask 比 outlines CPU bool 张量快。 (c) **多 schema 频繁切换**：xgrammar 编译快意味着新 schema 上线即用。 outlines 强在 **多字面常量 + 大量 enum** 的简单 JSON——jump-ahead 跳过的 token 数最多。

为什么受约束输出有时反而比无约束更快？

答案

两个原因： (a) **Jump-ahead 一次吐多 token**：受约束的"必出"片段（如 `"name":` 这类）一次 forward 吐 5-10 token，跳过 5-10 步 sample。 (b) **EOS 更早**：受约束输出严格终止于 schema 结束，不会像无约束那样"啰嗦多说几句"。实测复杂 JSON Schema 输出有时 TPOT 比无约束低 30%；简单 schema 仅打平或略慢（编译开销）。

13. 下一步

05-token-attention.md — 内存粒度的另一个核心议题。
09-advanced-features/02-structured-output.md — 进阶：复杂 schema 调优。
源码：sglang/python/sglang/srt/constrained/。
参考：LMSYS Blog "Compressed FSM"。

上游源码：sglang/python/sglang/srt/constrained/outlines_jump_forward.py。