AI 基建
0%
第五部分 · 推断与服务 · 第 35 章

结构化与长上下文推断

作者Changkun Ou
阅读时长约 13 分钟

有两种不改模型权重的服务时干预会改变解码器。结构化生成限制下一个词元的可选范围。长上下文策略限制注意力保留或读取哪些已经计算好的键值状态。两者共用同一个推断循环,提供的保证却不同。语法可以保证输出属于实现的形式语言;缓存驱逐与稀疏读取是对完整注意力的近似,必须结合具体任务评估。

这项区别在生产环境中很重要。可解析的 JSON 仍可能填错值。有界缓存也许能让流式输出保持通顺,却忘掉数千个词元之前出现的事实。两种机制都不能当作普遍的正确性保证。

约束解码会改变词元分布

来看 约束解码(constrained decoding),也就是带语法检查的解码循环。设 p(vx<t)p(v\mid x_{<t}) 为前缀 x<tx_{<t} 之后词表词元 vv 的模型概率,sts_t 为当前解析器配置,A(st)A(s_t) 为仍能导向有效完整结果的词元集合。语法约束采样使用

pG(vx<t,st)=p(vx<t)1[vA(st)]uA(st)p(ux<t).p_G(v\mid x_{<t},s_t)= \frac{p(v\mid x_{<t})\,\mathbf 1[v\in A(s_t)]} {\sum_{u\in A(s_t)}p(u\mid x_{<t})}.

其中,允许词元的 1[vA(st)]\mathbf 1[v\in A(s_t)] 为一,其他词元为零;分母将剩余的概率质量重新归一化。采样得到 vv 后,解析器消费该词元,并从 sts_t 推进到 st+1s_{t+1}。结束序列词元只有在解析器处于接受配置时才合法。如果允许集合为空,引擎就遇到了实现缺陷、不受支持的约束或不兼容的前缀。此时必须明确失败,不能退回原始分布继续采样。

A(st)A(s_t) 的定义比“下一个字符合法”更严格。词元只有在消费后至少仍有一个可接受的完整结果可达,才属于这个集合。否则,局部合法的词元可能把解析器带入死路,再也无法形成完整文档。若生成在接受配置中正常结束,正确实现能够保证输出字节串属于所实现的语言;最大词元数停止、取消或引擎故障仍可能留下不完整文档。

下面的可运行示例只演示重新归一化。它故意使用很小的允许集合,并不假装实现完整的 JSON 解析器。

from math import exp

def constrained_softmax(logits, allowed):
    if not allowed:
        raise ValueError("empty allowed set")
    peak = max(logits[i] for i in allowed)
    weights = [exp(logits[i] - peak) if i in allowed else 0.0
               for i in range(len(logits))]
    total = sum(weights)
    return [weight / total for weight in weights]

vocab = ["{", "}", '"', "name", "age", ":", ",", "true", "hello", "42"]
logits = [2.1, 0.3, 3.0, 1.2, 1.0, 0.5, 0.4, 0.8, 2.5, 1.1]
allowed = {2}  # Just after "{", this toy state accepts only an opening quote.
probabilities = constrained_softmax(logits, allowed)

for token, probability in zip(vocab, probabilities):
    print(f"{token:>6}  {probability:.3f}")
print("legal mass:", round(sum(probabilities[i] for i in allowed), 3))
print("illegal mass:", round(sum(probabilities[i] for i in range(len(vocab))
                                 if i not in allowed), 3))

解析器处理词元字节,而不是词元标签

模型输出分词器 ID,语法却定义在字节或 Unicode 码点等字母表上,运行时必须把两者组合起来。对于每个候选词元,运行时取得分词器产生的完整字节串,再判断解析器能否完整消费它。一个词元可能跨过多次语法转移,一个字符的字节序列也可能拆到多个词元中。特殊词元不一定会解码成普通文本,因此需要明确规则。

正则语言可以由 有限状态机(FSM),也就是记录部分输出当前所处语法状态的机器来识别。上下文无关文法还需要解析器栈来表示嵌套结构。以 JSON 对象为例,一般情形下它需要支持无限嵌套。若把所有约束都称为有限状态机,就会掩盖这种差异,并导致错误的缓存假设。

因此,生成循环中的五类状态必须彼此一致:

P 解析器配置 状态与栈 A 允许的词元 ID 仍有可接受的后续 P->A T 分词器映射 词元 ID → 完整字节串 T->A M 带掩码的采样器 归一化并选择 A->M U 提交词元 更新 KV 与解析器 M->U U->P 下一步
图 35.1. 一次结构化解码。解析器与分词器共同确定允许的词元 ID,采样器从中选择词元,采样结果同时更新模型状态与解析器状态。

保证覆盖什么

语法有效不等于取值正确。语法可以保证输出能解析为 JSON,也可以保证某个字段采用整数写法;它无法证明整数确实抄自原文、工具调用已获授权,或 SQL 语句可以安全执行。这些属于语义与策略检查。

JSON Schema 还有一层边界。引擎通常把自己支持的 JSON Schema 子集转换成语法或解析器。跨字段关系、远程引用、数值范围或字符串语义等关键字可能不受支持,也可能只在生成后检查。因此,实际承诺的保证是所请求 schema、后端语法覆盖、分词器集成与终止行为的交集。结构化生成引擎的基准测试显示,不同引擎在 schema 覆盖、编译行为和输出质量上存在明显差异 (Geng et al. 2025)。

消费方仍要再次验证完成的对象,拒绝额外字节、不受支持的 schema 特性、错误字段值、不安全操作和截断输出。约束解码消除了一类故障,不能取代应用验证器或授权层。

硬词元掩码也会改变模型的序列分布。上面的局部归一化通常不等于从原模型中按“最终属于语法”这一条件采样,因为它没有考虑每个前缀之后还剩多少有效概率质量。Grammar-Aligned Decoding 形式化了这项差别,并报告了普通语法约束解码虽然生成有效输出、内容质量却更低的案例 (Park et al. 2024)。影响取决于工作负载和模型,因此内容准确率与解析成功率都要比较。

控制语法执行成本

最直接的实现会在每个解码步逐一检查词表中的所有词元能否被当前解析器接受。这样每生成一个词元,都要重复数万次分词器与解析器工作。实用引擎会尽量把这部分工作移出解码循环。

对于正则约束,运行时可以预先计算每个自动机状态经过每个词表词元后的转移。Willard 和 Louf 提出一种词表索引,使“状态到允许词元”的查询成本很低 (Willard and Louf 2023)。编译缓存键必须包含规范化语法、分词器版本、特殊词元策略、字节编码与后端版本。若索引来自另一版分词器,就可能放行错误的词元 ID。

预计算不会让整个操作变成固定成本。掩码应用仍要访问选中的 logits 或词表大小的掩码;设备传输、同步、同一批次中的异构语法以及解析器栈更新,也可能成为可见开销。XGrammar 会预先检查不依赖上下文的词元,只用较小的上下文相关集合查询持久解析器栈,并让 CPU 上的语法工作与加速器执行重叠,从而降低上下文无关文法的开销。它在所评估引擎与工作负载上接近零的端到端开销是一项实测结果,不是普遍成立的上界 (Dong et al. 2025)。

强制跨度需要一次扩展计算

语法中常有标点和字段名等确定跨度。SGLang 会压缩只有一条出路的路径,直接跳到下一个分支 (Zheng et al. 2024)。这样可以避免每个强制词元各占一次自回归解码迭代,但这些词元并非免费。强制跨度仍必须写入模型状态,后续词元才能注意到它。运行时需要重新对边界分词,协调可能变化的分词结果,再以类似预填充的扩展计算一次处理确定词元,并扩展 KV 缓存。

B1 模型在语法 分支处采样 F 解析器找到一段 确定的字节跨度 B1->F R 重新对边界分词 并预留 KV 空间 F->R E 一次扩展/预填充计算 提交强制词元 R->E B2 在下一个分支 恢复解码 E->B2
图 35.2. 跳跃式处理。模型在一个语法分支采样后遇到确定的字节跨度,运行时重新对边界分词,并为整段内容扩展模型状态,然后才恢复采样。

这条路径与其他缓存更新一样,需要遵守预留、提交与回滚规则。如果边界重新分词改变了已经缓存的后缀,运行时必须让受影响的 KV 项失效并重新计算。扩展失败时不能让解析器状态领先于模型状态。跳跃式处理的性能取决于确定跨度长度、重新分词成本、扩展内核效率和调度器负载。

长上下文有四项彼此独立的限制

“支持 128K”无法回答所有长上下文问题。部署时必须区分四项限制:

  1. 模型支持的上下文长度。 位置编码、训练分布和架构决定模型在某个逻辑位置能否正常工作。驱逐 KV 项并不会扩展模型学到的能力。
  2. KV 容量。 常驻的键和值会随保留词元数、层数、头数、批量和元素宽度增长,因此可能限制系统能够接纳的并发量。
  3. 注意力流量。 即使完整缓存能够装下,每个新查询在解码时仍可能读取越来越多的状态。第 34 章的 IO 模型可以直接用于这里。
  4. 信息保留。 丢弃或跳过状态的策略可能删掉后续查询需要的证据。稳定的困惑度或流畅生成不能证明模型保留了长距离信息。

这里的 κ\kappa第 34 章 定义的每个保留词元的逻辑 KV 字节数。完整注意力为长度为 TT 的序列保留约

Mfull(T)=TκM_{\mathrm{full}}(T)=T\kappa

字节,尚未计入分配器与元数据开销。这里的 TT 是逻辑序列长度。若策略保留开头的 ss 个注意力汇词元和最近的 ww 个词元,则预算为 K=s+wK=s+w;预热结束后,最多保留

MwindowKκ+Mpolicy.M_{\mathrm{window}}\lesssim K\kappa+M_{\mathrm{policy}}.

其中,MpolicyM_{\mathrm{policy}} 包括块表、分数、页面摘要和填充。这条公式描述常驻的解码状态。某种方法即使会在完整处理提示词后压缩缓存,仍可能为全部 TT 个词元支付提示词预填充的峰值内存与计算。

三类策略,三种失效方式

比较长上下文方法时,按存储什么、读取什么来比较,比按论文发表时间更有用。

策略 建立后存储的 KV 每次查询读取的 KV 选择时机 主要失效方式
注意力汇加近期窗口 由注意力汇和窗口限定 全部保留项 按位置固定 被驱逐的中段证据无法回忆
H2O 式驱逐 由近期词元和历史重击者限定 全部保留项 生成过程中持续更新 历史注意力不一定能预测未来重要性
SnapKV 提示词压缩 压缩后的提示词状态加生成状态 全部保留项 提示词末尾的观察窗口 每个头未选中的提示词证据会被丢弃
查询感知的页面选择 完整缓存加页面摘要 选中的页面 每次查询重新选择 相关页面可能未进入前 kk 个读取集合

前三种策略通过删除条目来减少常驻 KV。除非系统保存了足够的源状态来重新计算,否则驱逐不可逆。查询感知的页面选择则保留完整缓存,只读取其中一部分,以降低注意力流量。下一次查询可以改选其他页面,所以当前未选中的页面仍可恢复,尽管本次注意力结果依然只是近似。

注意力汇与近期窗口

在一些已评估的 Transformer 系列中,存在 注意力汇(attention sink),即开头词元即使没有明确语义作用也持续吸引注意力的现象。StreamingLLM 报告称,保留少量初始 KV 项和一个近期窗口,可以避免朴素窗口丢掉前缀时出现的急剧退化。在论文评估的 Llama-2、MPT、Falcon 和 Pythia 配置中,该策略无需微调,就能让超长序列的流式语言模型评估保持稳定 (Xiao et al. 2024)。

论文把 softmax 归一化作为一种解释:当注意力头找不到明显相关的词元时,初始位置可以吸收部分注意力质量。这是观察到的模型行为,并非每个注意力头都必然遵守的定理。注意力汇数量和窗口大小取决于模型与工作负载。更重要的是,该策略不会记住已驱逐内容。它支持持续生成,不等于支持无限回忆。

图 35.3. 一个示意性的因果注意力模式,并非训练模型的实测结果。行表示查询,列表示键。最前面的列吸收注意力汇质量,近期键获得局部注意力,未来键的权重为零。

按注意力历史驱逐与提示词压缩

在 H2O 所评估的模型与任务中,它会保留近期词元以及累计注意力分数较高的词元 (Zhang et al. 2023)。这个分数来自历史。一个至今很少被注意的词元,仍可能对后续查询至关重要,因此重击者预算必须用未来才会用到信息的模式来测试,不能只看平均语言模型损失。

SnapKV 做的是另一件事。它观察提示词末尾附近的注意力,按注意力头分别选择重要的提示词位置,并在生成开始前压缩提示词 KV (Li et al. 2024)。这是提示词压缩,不是持续更新的重击者缓存。它可以减少解码状态的内存与读取量,却未必降低预填充峰值成本,因为选择发生在提示词已经处理完之后。

查询感知的页面选择

Quest 保留完整缓存,并为每个页面保存逐通道的键极值。对于新查询,这些摘要给出一个分数上界,供运行时选择注意力要加载的前 kk 个页面 (Tang et al. 2024)。存储量仍接近完整 KV 加摘要,但设备内存流量与注意力计算可以减少。这让失效方式从永久删除变成了单次查询遗漏,却不能保证每个相关页面都被选中。

训练原生稀疏注意力的架构,会把词元选择纳入训练。它们是重要的替代方案,却不属于上面这些固定模型推断策略。给稠密注意力检查点事后加驱逐或选择规则,与训练模型使用稀疏注意力,是两种不同干预。

缓存状态必须始终自洽

保留或读取策略操作的不只是键和值数组。每个保留项还需要逻辑位置、层与头的身份、数值格式和所有权状态。位置元数据必须与模型的旋转位置编码或其他位置编码保持一致。除非方法明确要求并验证这种转换,否则把有缺口的保留词元重新连续编号会改变注意力几何。

当策略按完整块保留或驱逐时,分页分配确实有帮助;但分页分配并不会让任意逐词元或逐头驱逐变得免费。一个物理块可能同时包含保留与丢弃槽位,可能被共享前缀缓存的请求共同持有,也可能还有读取者尚未结束。细粒度策略需要按块选择、间接 gather、压实或专用稀疏内核。来自 第 32 章 的引用计数、预留、提交、回滚、取消与前缀缓存身份规则仍然适用。

策略元数据同样占用内存和带宽。峰值与常驻测量必须包括 H2O 式分数、SnapKV 的逐头索引和 Quest 的页面摘要。某个策略即使减少了 KV 存储,如果启动了缓慢的 gather 内核,也可能出现容量改善了,延迟仍可能变差的结果。

验证完整服务路径

结构化生成与长上下文策略需要不同的正确性测试,之后还要接受同一套负载测试。

对于结构化生成:

  1. 用官方解析器或 schema 验证器检查对抗性和真实 schema。覆盖嵌套对象、Unicode、跨越多次转移的词元片段、特殊词元、空允许集合与达到最大长度时的停止。
  2. 分别报告 schema 有效率、语法覆盖率、字段级正确率与下游拒绝率。字段填错的有效对象,不算成功的抽取或工具调用。
  3. 测量语法编译延迟、编译缓存命中率与内存、掩码准备和应用时间、允许集合大小以及每输出词元耗时。还要测试混合语法批次。
  4. 在不同分词边界、取消和分配失败条件下,把跳跃式处理的输出与 KV 状态同普通逐词元扩展比较。

对于长上下文推断:

  1. 在模型支持的上下文长度内与完整 KV 基线比较。扫描提示词长度、生成长度、缓存预算与页面选择预算。
  2. 按证据位置报告检索准确率,并报告多跳回答质量、长生成质量,以及“早期注意力很低的事实后来变得必要”等对抗案例。困惑度本身不够。比较研究显示,缓存压缩策略的排名会随任务与模型变化 (Yuan et al. 2024)。
  3. 在同等接纳负载下测量提示词峰值内存、常驻与峰值 KV 内存、元数据、每步读取的 KV 字节、首词元时间、每输出词元时间、吞吐量、有效吞吐量与延迟尾部。
  4. 测试共享前缀、块引用计数、取消、位置边界、量化 KV、不受支持的内核,以及压实或选择失败后的回滚。
争议所在

不存在脱离工作负载的最佳缓存策略。流式生成也许能容忍注意力汇加窗口预算,问答任务却可能因为一句被驱逐的话而失败。基于注意力历史的分数在一种模型上可能有效,在另一种模型上却会漏掉未来依赖。查询感知选择虽然保留已存储的 KV,仍然近似了每次查询读取哪些页面。每项质量或速度声明都必须同时报告模型、任务、提示词分布、生成长度、预算、内核和完整缓存基线。

下层约束

前几章决定这些策略能否以足够低的成本运行。调度器必须容纳语法编译与异构掩码;分配器必须在驱逐或压实过程中保持缓存所有权;量化契约必须覆盖所有保留的 KV 与策略元数据;投机解码必须把同一语法应用于目标分布。下一章还会加入多模态词元,它们更长且更不均匀的前缀会进一步加重语法批处理、预填充峰值和缓存选择的压力。

收益与边界

当语法、分词器与终止行为彼此一致时,约束解码可以让语法错误变得不可达,却无法让取值变成事实。长上下文策略可以限制常驻状态,或减少注意力读取,却无法让已经丢弃的证据重新可用。达到生产可用,意味着先明确这条边界,再把保证、任务质量与服务结果一起衡量。

延伸阅读

  • Willard & Louf, “Efficient Guided Generation for Large Language Models,” 2023. arXiv:2307.09702
    Willard 与 Louf 将形式语言和模型词表组合起来,使正则表达式与文法约束能在生成时高效提供词元掩码。
  • Geng et al., “JSONSchemaBench: A Rigorous Benchmark of Structured Outputs for Language Models,” 2025. arXiv:2501.10868
    JSONSchemaBench 以真实模式和官方模式测试,从效率、约束覆盖与输出质量等方面评估结构化生成系统。
  • Park et al., “Grammar-Aligned Decoding,” 2024. proceedings.neurips.cc
    Grammar-Aligned Decoding 表明,普通的局部文法掩码即使保证语法有效,也可能扭曲模型的序列分布。
  • Dong et al., “XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models” (用于结构化语法的文法约束解码), 2025. proceedings.mlsys.org
    XGrammar 通过预检查词元、持久化解析栈,以及文法工作与加速器执行重叠来加速上下文无关文法。
  • Xiao et al., “Efficient Streaming Language Models with Attention Sinks,” 2024. arXiv:2309.17453
    StreamingLLM 在近期窗口之外保留少量初始注意力汇词元,使论文评估的模型族保持稳定的流式语言建模。
  • Zhang et al., “H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models,” 2023. proceedings.neurips.cc
    H2O 是一种动态 KV 缓存淘汰策略,在论文评估的模型与任务中平衡近期词元和按累积注意力识别的重击者。
  • Li et al., “SnapKV: LLM Knows What You Are Looking for Before Generation,” 2024. proceedings.neurips.cc
    SnapKV 使用提示末尾的观察窗口,按头选择并聚合重要的提示 KV 位置,再开始生成。
  • Tang et al., “QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference,” 2024. proceedings.mlr.press
    Quest 保存逐页键极值,并用每次查询选择注意力读取的前 K 个 KV 页,同时保留完整缓存。
  • Yuan et al., “KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches,” 2024. aclanthology.org
    该基准在七类任务上比较长上下文效率方法,发现质量与效率权衡会随方法、任务和模型而变化。

评论

登录后评论