跳到正文
LLM 推理
S15解码加速·338

结构化输出

合法 JSON 不是一个 prompt 问题。把 schema 编译成状态机,然后把语法不允许的 logit 全部屏蔽掉。

  • 有限状态机 / 下推自动机
  • logit 掩码
  • token 级语法
  • 跳跃解码

为什么重要

问题

你需要模型返回符合某个 schema 的 JSON,因为下游有程序要解析它。通常的做法是在 prompt 里好言相劝,失败了就重试。它在大多数时候有效,而对生产数据链路上的东西来说,「大多数时候有效」几乎是最糟糕的性质。

重试要浪费一整次生成,而且不收敛:一个对 schema 犯迷糊的模型会继续犯迷糊。失败模式还是无声的。少一个字段照样能解析,然后在下游第三个服务里把什么东西搞坏。

核心思路

解法

别再求它了。采样器本来就要决定下一个 token 是什么,而你本来就知道此刻语法允许哪些 token。把两者取交集:在 softmax 之前把所有非法 token 的 logits 设为 −∞,模型就从机制上不可能产出非法输出。

受约束的采样python
def constrained_sample(logits, fsm_state, fsm, tokenizer, params):
    # Which token ids can legally follow, given where the automaton is?
    allowed = fsm.allowed_token_ids(fsm_state)          # precomputed per state

    mask = np.full_like(logits, -np.inf)
    mask[allowed] = 0.0
    logits = logits + mask                              # that is the whole mechanism

    token = sample(logits, **params)
    return token, fsm.advance(fsm_state, token)

有两点值得注意。模型的偏好在合法 token 之间得以保留:你只是移除了选项,并没有推翻它的判断。而代价只是一次向量加法,相对于产生这些 logits 的那次前向来说,等于免费。

语法定义在字符上,采样发生在 token 上

自动机反倒是容易的部分。难在像 ": {" 这样一个 token 可以同时携带一个字符串终止符、一个冒号、一个空格和一个对象起始括号,一次跨越四个状态转移,因此掩码不能按字符来做。你必须为每个自动机状态预先算出:哪些 token id 的完整字符展开能让自动机始终停留在合法状态。Outlines 和 XGrammar 这类库存在的意义,就是做这件预计算。

工作原理

工作原理

示意图从 schema 到自动机再到 logit 掩码
编译一次,每步都做掩码JSON schema或正则,或 EBNF自动机FSM 或下推自动机token 级掩码每个状态预先算好logits += mask非法项置 −∞采样永远合法自动机的一个片段期待 {期待 "期待 key期待 :期待 value, or }下一个 key掩码对 LOGITS 做了什么namekeptagekeptactivekeptcat−∞{−∞7−∞true−∞在合法 token 之间,模型的偏好被完整保留:掩码改变的是「什么可能」,不是「更偏好什么」。跳跃前进解码当只有唯一一个 token 合法时,干脆别跑模型。结构化输出可以比自由输出更快。TOKENIZER 带来的麻烦语法是定义在字符上的;而采样是在 token 上的。一个 token 可能横跨一次状态转移,所以掩码必须按 token 算,不能按字符算。
编译对每个 schema 只做一次并被缓存。生成时的代价只是一次状态查表和一次向量加法。

正则不够用:JSON 需要一个栈

有限状态机能表达正则语言,用来描述电话号码、日期或固定枚举都够。JSON 不是正则的:嵌套的对象和数组需要任意深度的括号匹配,任何有限状态集合都追踪不了。

所以通用的语法约束解码用的是下推自动机:状态机加一个栈。XGrammar 和 llguidance 都这么做,并带有一个很关键的优化。大多数 token 不影响栈,于是它们把词表切成「上下文无关 token」(可用预计算位图检查)和一小撮「上下文相关」的剩余部分。

跳跃解码:真正的免费午餐

在只有一个必需 key 的 schema 里,{ 之后接下来的好几个字符是确定的:引号、key 名、引号、冒号。模型根本没有选择可做。

那就别问它。当自动机只允许一种后续时,把整段被强制的内容直接输出并推进,不需要任何采样决策。在真实引擎里,这些被强制的 token 仍然需要前向来填充 KV 缓存,但许多串行的 decode 步骤会坍缩成对整段内容的一次批量扩展,便宜得多。(本章的玩具模型不保存 KV 缓存,所以干脆跳过了前向。)在 schema 密集的输出上,这能省掉相当大一部分串行 decode 步骤。受约束生成因此常常比生成同样文本的无约束版本更快,而不是更慢。

它解决不了什么

受约束解码保证的是语法合法性,不是语义正确性。要求有 country 字段的 schema 一定会拿到一个;但没有任何东西逼它是正确的国家。而且过度约束会有害:如果模型想说明自己无法回答,而语法只允许一个完整的对象,你就逼出了一次自信的编造。

实用的缓解办法是在 schema 里包含一个「拒绝」形态,也就是与 {"error": string} 做联合类型,让模型有一条合法的路可以说不。

动手观察

动手试试

模拟器受语法约束的 JSON
第 0 步
正在强制的 schema
{ "name": string, "age": number, "active": boolean }
当前输出

(还没有输出)

状态机状态
期待对象起始括号
允许的 token
1/28
免费跳跃次数
0
能否解析为 JSON
尚未完成

期待对象起始括号. 只有高亮的 token 可以被采样,其余 token 的 logits 在 softmax 之前就被置为 −∞。非法输出不是被「不鼓励」,而是根本不可能

词表掩码 · 第 0 步
{}":,nameageactivetruefalsenull012379AdaLinKimReythecatsat[]-
  • 状态机允许
  • 被掩码为 −∞

跳跃解码:此刻只有一个 token 合法,于是引擎完全不跑前向就直接输出它。结构化输出可以比非结构化更快。

开着约束跑一遍。每个随机种子都产出符合 schema 的合法 JSON,而 ⚡ 标记显示了哪些步骤根本不需要前向。关掉约束再跑:采样器立刻开始乱走,因为没有任何东西拦着它。这里的「模型」是均匀随机的,拿来漫画式地刻画一个还没内化该 schema 的真实模型,倒也不算冤枉它。

亲手实现

实现

编译这一步把每个自动机状态映射到一个 token id 位图。它就是全部成本,也正是你要按 schema 缓存它的原因。

code/s15_structured_output.py(节选)python
def compile_masks(fsm, tokenizer):
    """For every state, which token ids keep us inside the grammar?"""
    masks = {}

    for state in fsm.states:
        allowed = []
        for token_id, token_str in tokenizer.vocab.items():
            # Feed the token's characters through the automaton one at a
            # time. A token is legal only if EVERY character is accepted.
            s = state
            if all((s := fsm.step(s, ch)) is not None for ch in token_str):
                allowed.append(token_id)
        masks[state] = np.array(allowed, dtype=np.int32)

    return masks     # cache this per (schema, tokenizer) pair

编译是 O(状态数 × 词表大小),而词表有 13 万

在一个大 schema 上朴素地编译可能要花好几秒;忘了缓存的话,就是每个请求都花好几秒。生产库会按 schema 哈希缓存、在状态被首次到达时惰性编译、并把掩码存成位集而不是下标列表。如果受约束解码莫名其妙地慢,先去看编译器,再去看前向。
在本地运行
把一个 JSON schema 编译成 token 级自动机,用一个玩具模型生成受约束输出,断言每个样本都能解析并通过 schema 校验,并测量跳跃解码跳过了多少 decode 步骤。
$ python code/s15_structured_output.py
预期输出: 200/200 的样本符合 schema;在这个 schema 上约有 62% 的 decode 步骤被跳跃解码省掉;以及一条无约束基线:即使模型每步有 99% 的概率选中合法 token,整个对象也只有约 84% 的概率有效。

只需要 NumPy — 查看环境准备.

生产实践

在生产环境中

  • Outlines —— 正则与 JSON schema 约束生成;把预编译 FSM 这条路带火了。
  • XGrammar —— 下推自动机,带上下文无关/相关 token 的切分;目前 vLLM 的默认后端。
  • llguidance —— Guidance 背后的后端,目标类似,Rust 实现很快。
  • SGLang —— 显式实现了跳跃解码,并在 schema 密集负载上报告了很大的加速。
  • picoLM —— 这个想法的极简版本,约 175 行 C:预分析词表并对 logits 做掩码以强制 JSON 结构,零依赖。

练习

  1. 1
    用一个显式的栈扩展自动机以支持嵌套对象,并验证不配对的括号变得不可达。
  2. 2
    实现跳跃解码,分别测量一个「宽」schema 和一个「深」schema 上被跳过的 decode 步骤比例。
  3. 3
    加一个联合类型,让模型可以输出 {"error": "..."} 而不是编造答案;然后构造一个 prompt,让无约束模型选择拒绝、而过度约束的模型开始胡编。

继续学习

接下来

在把引擎分布式化之前,还有最后一项架构技术。S16 讲专家混合推理:模型每个 token 只激活 5% 的参数,但整个模型仍然必须常驻显存。

自测

习题

先作答,再看解析。答错比答对更有价值,因为解析会指出你该回头重读哪一部分。

自测 1 题 / 共 6

为什么语法约束不能一个字符一个字符地施加?

得分 0/6