DeepInfer 源码完全拆解（六）：结构化输出与确定性——语法树异步编译、投机回滚与 batch-invariant

1.1 问题：自由文本采样与格式要求天然冲突

大模型每一步的输出，本质上是词表上的一个概率分布，采样就是在 $\lvert V \rvert$ 个候选里抽一个。可是真实业务里经常不是「随便说点什么」，而是「必须吐出一个合法的 JSON」：

• 函数调用（function calling）：模型输出的必须是符合某个 schema 的参数对象，下游要 json.loads 后直接使用；
• 数据抽取：把一段自然语言抽成固定字段的表格行；
• 代码生成：生成一段必须能通过语法检查的 SQL / 正则 / 表达式；
• 受限枚举：只能回答 yes / no，或者只能是形如 AB-123 的编号。

最朴素的方案是「事后校验 + 重试」：让模型自由生成，解析失败就重来。代价很大：重试意味着一次完整解码（几百个 token）被浪费，而且失败率高的时候延迟方差会爆炸。一个中等规模的线上服务里，几十毫秒的 TTFT 后面跟着的可能是几秒的重复解码。

结构化输出（structured output，也叫 constrained decoding / guided decoding）换了个思路：不要事后校验，而是在采样那一刻就把不合法的 token 置为不可能。既然每步的采样是从一个分布里抽，那么只要把分布中所有「会破坏格式」的候选概率压到 $-\infty$，采样就永远抽不到它们。这样生成的序列从构造上就是合法的，不需要重试，也不会出现「生成了一半才发现格式错了」。

1.2 技术核心：把格式变成一个「状态机」，每步只放行一部分 token

要让「哪些 token 合法」可计算，需要把格式编译成一个形式语言的接受器。三种常见写法：

无论哪种，编译结果都可以抽象成同一类对象：当前处于哪个状态，以及每个候选 token 会让状态迁移到哪里。于是每一步推理的流程就变成：

其中 bitmask 是一个长度为 $\lceil \lvert V \rvert / 32 \rceil$ 的 int32 数组，第 $i$ 位为 1 表示词表第 $i$ 个 token 在当前状态合法。下面这张图来自 vLLM 官方博客对结构化解码的图解，画的是「JSON schema → 状态机 → 逐位掩码」的直观过程：

（图源：vLLM 官方博客：Structured Decoding in vLLM）

这里有一个容易被忽视的关键点：状态机是逐 token 推进的，而 token 不等于字符。BPE 分词下，一个 token 可能包含多个字符，甚至跨过格式边界（例如 {"a": 可能是一个 token）。这意味着「当前状态 + 候选 token」的合法性判断必须建立在tokenizer 的实际词表上，而不是在字符层面做正则匹配。这正是为什么需要一个专门的库（XGrammar）来把 TokenTokenizerInfo 和文法编译结合起来——也是为什么下面会看到 TokenizerInfo.from_huggingface(tokenizer, vocab_size=...) 这样的调用。

1.3 XGrammar：让「每步算掩码」变得可接受

XGrammar 由 MLC 团队提出（论文 arXiv:2411.15100，XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models）。它要解决的核心矛盾是：语法状态机的逐 token 掩码计算必须极快，否则它会成为解码循环里新的瓶颈。

XGrammar 的做法分两层：

• 编译期：把语法拆成「上下文无关部分 + 词表相关的部分」。上下文无关的展开（比如 JSON 的括号嵌套结构）预编译成一张紧凑的状态转移表；词表相关的部分（哪些 token 对应哪些字符序列）单独索引。这一步把大部分工作从运行时挪到了编译时。
• 运行期：每个解码步只需要对当前状态查表，填出 bitmask。由于 bitmask 是稠密的位图，写入是纯内存操作，可以很好地与 GPU 侧的计算重叠。

XGrammar 的论文报告称，相比此前逐字符推进的约束方案，其掩码计算开销显著下降；vLLM 官方博客给出的实测对照里，用 XGrammar 之后每输出 token 的结构化开销被压到可忽略的量级：

（图源：vLLM 官方博客：Structured Decoding in vLLM）

支撑「快」的两条关键思路值得展开一点，因为它们解释了为什么掩码能便宜到可以放进解码循环：

其一，压缩有限状态机（compressed FSM）。 直接把正则/JSON schema 展开成原始 FSM 会得到一个字符级的巨大状态图；而真正在解码时使用的是token 级的转移。压缩的做法是：把「每个状态能接受哪些 token」预先算好并合并等价状态，让运行期只需查表。SGLang 团队 2024 年 2 月发布的博客 Compressed Finite State Machine 用图解把「原始 FSM → 压缩 FSM」的规模差异画得很直观：

（图源：SGLang 官方博客：Compressed Finite State Machine）

其二，自适应 token 掩码缓存。 并非所有状态都需要重新算掩码。XGrammar 把词表 token 分成「与上下文无关」与「与上下文相关」两类，对前者缓存判定结果：

（图源：SGLang 官方博客：Compressed Finite State Machine）

实测收益方面，该博客给出的结论是压缩后状态数与每次掩码生成的代价都大幅下降：

（图源：SGLang 官方博客：Compressed Finite State Machine）

这里需要说明一个工程前提：DeepInfer 直接复用 XGrammar 库，而不是自己实现文法引擎。仓库里 deepinfer/structure_output/ 这一层的职责不是「发明算法」，而是「把这个库塞进一个高并发、投机解码、多实例部署的引擎里，并且不让它成为正确性或性能的破口」。下一节开始逐文件拆这一层。

2.1 一个请求要过四道手

deepinfer/structure_output/ 目录下只有四个文件，职责切得很干净：

先看编译。compiler.py 里 dispatch_and_compile 按 schema_type 分派到 XGrammar 内部的三个入口，并且给每一类失败都补上上下文：

1
def dispatch_and_compile(self, schema_type: str, schema: str) -> CompiledGrammar:
2
    try:
3
        if schema_type == "json":
4
            ctx = self.inner_compiler.compile_json_schema(schema, strict_mode=True)
5
        elif schema_type == "regex":
6
            ctx = self.inner_compiler.compile_regex(schema)
7
        elif schema_type == "ebnf":
8
            ctx = self.inner_compiler.compile_grammar(schema)
9
        else:
10
            raise ValueError("Unsupported Grammar Type")
11
    except RuntimeError as e:
12
        raise RuntimeError(f"xgrammar compilation RE: {e}") from e
13
    except ValueError as e:
14
        raise RuntimeError(f"xgrammar invalid input: {e}") from e
15
    except Exception as e:
16
        raise RuntimeError(f"xgrammar unknown error: {e}") from e
17
    return ctx

三个分派分支值得说明为什么这么选：

• json 走 strict_mode=True。strict 模式下 schema 的语义被收紧（例如不允许未声明的额外字段、要求 required 全部出现）。对函数调用这类下游要直接消费的场景，宽松模式会导致模型生成一个「schema 之外但语法合法」的对象，下游反而更难处理；strict 让生成的文本与 schema 一一对应。
• regex 与 ebnf 分别走 compile_regex 与 compile_grammar。注意 ebnf 这一路对应的是上下文无关文法，比正则表达式的表达能力严格更强（能表达括号配对），这也是为什么 SQL 这类嵌套结构必须走 EBNF 而不能用正则。
• 三种 except 分支的理由不同：RuntimeError 通常是 XGrammar 内部编译失败（文法无解、状态爆炸），ValueError 是输入本身非法（比如传入的 schema 不是合法 JSON），其他异常归为未知。分成三类是为了让上层日志能区分「用户给了坏 schema」和「库内部炸了」——前者是用户问题，应该直接判死该请求；后者需要运维介入。

2.2 编译是多慢？慢到不能同步等

编译一个复杂 JSON Schema 要做状态展开，是 CPU 上几十毫秒到几百毫秒量级的操作。如果放在请求入口同步执行，那么这段时间里：

• GPU 可能在等（如果整个 batch 都被这个请求卡住）；
• 更糟的是，在大批量场景下，一个请求的编译会阻塞整个调度循环——因为 add_request 是在引擎主线程里跑的。

DeepInfer 的做法是把 matcher 的构建整个挪到一个线程池里，so_matcher_builder.py 的 SOMatcherBuilder：

1
self._pool = ThreadPoolExecutor(max_workers=max_workers, initializer=initializer)
2
self._compiler = self._pool.submit(XGrammarCompiler, self._tokenizer_info, num_threads=2).result()
3
# submit/pump 均在主线程调用，worker 线程只填 future，不碰 seq 入调度逻辑
4
self._pending: list[tuple[object, XGrammarMatcherFuture, Callable[[object], None]]] = []

三个设计点：

第一，编译器本身也是异步建的。 XGrammarCompiler 的构造被 submit(...).result() 包起来——注意这是在 SOMatcherBuilder.__init__ 里等待的，属于一次性初始化成本。之所以不直接同步 XGrammarCompiler(...)，是因为它内部要建 TokenizerInfo、初始化 cache 结构，在 bind_cpus 场景下还需要先设置 CPU 亲和性；放进线程池是为了复用同一套 initializer 逻辑（os.sched_setaffinity(0, set(bind_cpus))）。

第二，CPU 亲和性绑定的是 worker 线程，不是进程。 代码注释很明确：

1
def _bind_cpu_affinity() -> None:
2
    # 0 = 调用线程自身：绑定执行 initializer 的 worker 线程，而非整个进程/主线程。
3
    # psutil.Process() 的默认 PID 指向主线程，不能用于绑定当前 worker 线程。
4
    os.sched_setaffinity(0, set(bind_cpus))

这是一个真实的坑：如果按常规写法用 psutil.Process().cpu_affinity([...])，绑定的是主线程的 PID，而 initializer 是在 worker 线程里跑的，结果是「亲和性设置没生效还占用了主线程的核」。用 os.sched_setaffinity(0, ...) 里的 0 表示调用线程自身，才符合意图。为什么需要绑核？因为编译是纯 CPU 密集的，如果它和调度主线程、tokenizer 线程抢同一批核，编译期间解码步的延迟就会出现尖刺。DEEPINFER_SO_BUILDER_CPUS（默认 4）就是这个核数的配置项。

第三，worker 线程不碰调度。 注释写着「submit/pump 均在主线程调用，worker 线程只填 future，不碰 seq 入调度逻辑」。这是并发正确性的边界：seq 是引擎的共享对象，入队、状态迁移必须在主线程完成，否则会与调度器的状态机打架。线程池只负责「生产一个 matcher 对象」，对象的所有权转移通过 pump 在主线程序列化地完成。

2.3 submit / pump：把「还没编译完」变成一个正常的中间状态

异步之后必然出现一种状态：请求已经到达，但它的 matcher 还没编译完。DeepInfer 用 submit + pump 一对函数处理这个窗口：

1
def submit(self, seq, so_args, on_ready, *, replay_tokens=None) -> None:
2
    """异步构建 matcher，编完后由 pump 挂到 seq 上并调用 on_ready(seq) 入调度。"""
3
    future = self.build(so_args, replay_tokens=replay_tokens)
4
    self._pending.append((seq, future, on_ready))
5

6
def pump(self) -> None:
7
    """轮询 pending，把已编完的 seq 出队并 on_ready 入调度。"""
8
    if not self._pending:
9
        return
10
    still_pending = []
11
    for seq, future, on_ready in self._pending:
12
        matcher = future.try_result()  # 非阻塞，超时返回 aborted_matcher
13
        if matcher is None:
14
            still_pending.append((seq, future, on_ready))
15
            continue
16
        seq.so_matcher = matcher
17
        on_ready(seq)
18
    self._pending = still_pending

pump 在每个引擎步开头被调用一次，try_result 是非阻塞的：编译完了就把 matcher 挂到 seq.so_matcher 上并调 on_ready(seq)（也就是真正把请求放进调度器队列）；没编完就留在 _pending 里等下一轮。

这个设计有几个直接收益：

• 编译与解码并行。请求 A 的 schema 正在编，同一步里请求 B、C 的解码照常进行。编译是 CPU 活，解码是 GPU 活，两者天然可以重叠。
• 入队顺序 = 编译完成顺序。先编完的先被调度，这是一个自然的背压形式——复杂 schema 慢一点被接纳，而不是让它的编译时间污染整个 batch 的开始时间。
• has_pending 参与可调度判定。引擎用 self._so_matcher_builder.has_pending() 和 get_num_pending() 把「还在编译中的请求数」算进「引擎是否还有活」的判断里，避免引擎误以为队列空了就开始排空/退出。

2.4 超时与失败：aborted_matcher 是统一降级口

异步就一定会有「编译太久」或者「编译炸了」的情况。XGrammarMatcherFuture 处理这个，超时时间来自 DEEPINFER_XGRAMMAR_TIMEOUT（默认 20 秒）：

1
def result(self) -> XGrammarMatcher:
2
    ...
3
    rem_time = self.timeout - (time.monotonic() - self.start_time)
4
    rem_time = max(rem_time, 1e-5)
5
    try:
6
        self._resolve(rem_time)
7
    except concurrent.futures.TimeoutError:
8
        self.res = XGrammarMatcher.aborted_matcher("xgrammar timeout when waiting for result")
9
    except RuntimeError:
10
        self.res = XGrammarMatcher.aborted_matcher("xgrammar runtime error when waiting for result")
11
    except Exception as e:
12
        self.res = XGrammarMatcher.aborted_matcher(f"xgrammar unknown Error: {e}")
13
    self._log_abort()
14
    return self.res

注意失败不是抛异常，而是返回一个 aborted_matcher——一个 ctx=None、aborted_msg 非空的特殊 matcher。这是整个 SO 层最重要的工程决定：

编译失败、超时、运行期异常，一律收敛成「这个请求带了一个 aborted 标记」，由上层用统一的判死路径处理，而不是让异常在调度主循环里炸开。

为什么必须这样？因为在 DeepInfer 里，采样发生在所有 rank 同步的路径上（CP rank 0 采样，其他 rank 参与集合通信）。如果某个请求的 matcher 在采样主路径抛出 Python 异常，会连坐同批所有请求，并且导致 NCCL barrier 超时——一个用户的坏 schema 拖垮整个引擎。代码注释把这个动机写得很直白：

1
except Exception:
2
    # 此处在 CP rank 0 采样主路径，raise 会连坐所有请求 + NCCL barrier 超时；
3
    # SO 的 bug 只判死本请求，不带崩 engine。

aborted_matcher 的构造也很有意思，它是一个完全空的 matcher：

1
@staticmethod
2
def aborted_matcher(msg: str) -> XGrammarMatcher:
3
    ret = XGrammarMatcher(None, None, None, [])
4
    ret.aborted_msg = msg
5
    return ret

ctx=None 意味着 inner_matcher=None；release() 会把所有成员清空，__del__ 里 suppress(Exception) 保证析构不抛。所以一个 aborted matcher 在生命周期上是安全的：它不会被推进，只会被识别并判死。

2.5 请求入口的两种构建模式

回到 llm_engine.py 的 add_request 路径，能看到异步与同步并存：

1
so_pending = False
2
if sampling_params.so_args is not None:
3
    so_args = sampling_params.so_args
4
    if not enqueue:
5
        # 调用方自行管理调度，无法靠 pump 重入，保持同步构建
6
        seq.so_matcher = self.so_matcher_builder.build(so_args).result()
7
    else:
8
        so_pending = True
9
...
10
if enqueue:
11
    def _enqueue(s: Sequence) -> None:
12
        if s.so_matcher is not None and s.so_matcher.is_aborted():
13
            s.status = SequenceStatus.FINISHED_SO_FAILED
14
        self.scheduler.add_sequence(s)
15
        ...
16
    if so_pending:
17
        self.so_matcher_builder.submit(seq, sampling_params.so_args, _enqueue)
18
    else:
19
        _enqueue(seq)

• enqueue=True（常规路径）：异步提交，_enqueue 作为回调，等编译完再入队。
• enqueue=False（调用方自己管调度，比如某些离线批处理工具）：走同步构建，理由是「无法靠 pump 重入」——没有引擎主循环就没有 pump 的机会窗口，只能同步等。

_enqueue 里第一句就是检查 is_aborted() 并直接置 FINISHED_SO_FAILED。这里定义了一个专门的终态枚举值（而不是复用一个通用的 FINISHED_ABORTED），目的就是让监控能区分「用户主动取消」与「schema 编译失败」。

2.6 两个安全细节：缓存上限与「离线不接受序列化文法」

compiler.py 里有一段关于缓存的日志，值得单独拉出来讲：

1
self.disable_cache = envs.DEEPINFER_XGRAMMAR_DISABLE_CACHE
2
self.cache_limit_bytes = envs.DEEPINFER_XGRAMMAR_CACHE_LIMIT_BYTES
3

4
if self.disable_cache:
5
    logger.info("Cache for xgrammar is disabled.")
6
elif self.cache_limit_bytes <= 0:
7
    logger.warning(
8
        "xgrammar cache is enabled WITHOUT a size limit "
9
        f"(DEEPINFER_XGRAMMAR_CACHE_LIMIT_BYTES=%d). "
10
        "The cache grows unboundedly with distinct schemas and may cause OOM (DoS vector, "
11
        "see CVE-2025-32381). Set a positive byte limit for production."
12
    )
13
else:
14
    logger.info(f"The cache size limit for xgrammar is set to {self.cache_limit_bytes // 1024 // 1024}MB.")

XGrammar 会缓存编译好的 grammar，以便同一个 schema 被反复使用时跳过编译。但缓存如果不设上限，就是一个内存耗尽攻击面：攻击者不断发送互不相同的 schema，缓存条目数随请求量线性增长，最终 OOM。这正是 CVE-2025-32381（XGrammar 编译缓存无上限导致拒绝服务），影响版本为 0.1.18 之前。DeepInfer 的处理是：

• 默认打开缓存但设置上限，DEEPINFER_XGRAMMAR_CACHE_LIMIT_BYTES 默认 10 GiB；
• 上限被配成 $\le 0$ 时不是静默继续，而是打一条带 CVE 编号的 warning；
• 再留一个 DEEPINFER_XGRAMMAR_DISABLE_CACHE 作为 kill switch，注释写明是「xgrammar <= 0.1.x had production incidents with cache enabled」。

最后，build() 里有两句显式的拒绝，是接口演进留下的护栏：

1
if "serialized" in so_args:
2
    raise ValueError("Offline structured output does not accept serialized grammars; provide schema_type and schema")
3
if "bound_token_ids" in so_args:
4
    raise ValueError("bound_token_ids is no longer supported; use bound_token_id with a single non-negative integer")

第一句的原因很实际：序列化过的 grammar 是跨版本、跨平台的二进制，离线场景下无法保证本机编译器和它兼容，宁可要求用户传 schema_type + schema 让本机重新编译。第二句是从复数 bound_token_ids 收敛到单数 bound_token_id 的接口清理——顺便说，旧写法在这里是报错而不是被兼容处理，属于「宁可让用户改代码，也不留两套语义」的风格。

2.7 状态恢复通道：replay_tokens 与 copy_without_status

build() 里还有一段容易被忽略的逻辑：

1
def _build() -> XGrammarMatcher:
2
    ctx = self._compiler.dispatch_and_compile(so_args["schema_type"], schema)
3
    matcher = XGrammarMatcher.from_compiled(ctx, self._vocab_size, self._eos_id, trigger_tokens, bound_token_id=bound_token_id)
4
    for i, token in enumerate(replay_tokens):
5
        if not matcher.accept_token(int(token)):
6
            raise RuntimeError(f"SO replay failed: token {token} (idx {i}/{len(replay_tokens)}) rejected during matcher state restore")
7
    return matcher

replay_tokens 是一条状态恢复通道：新 matcher 建好之后，把一批历史 token 重新喂进去，让状态机回到「已经消费过这些 token」的位置。为什么要这个？因为 matcher 是「随请求生灭」的 Python 对象，而引擎里存在几种会让请求带着已有历史重新开始的场景：KV 前缀缓存命中的续写、序列 fork 复制前缀、离线工具把一条长序列拆成多段推进。这些场景下如果 matcher 从零状态出发，语法状态就会与实际文本位置错位。

回放失败的处理也值得注意：不是忽略，而是直接 RuntimeError 让整个 matcher 构建失败，走 aborted 路径判死请求。理由很硬：回放失败意味着「历史 token 序列本身不满足这个语法」，此时无论怎么继续都是错的，不如尽早失败。

与之配对的是 matcher 侧的 copy_without_status()：

1
def copy_without_status(self) -> XGrammarMatcher:
2
    ret = XGrammarMatcher(self.ctx, self.eos_token_id, self.mask_after_eos, self.trigger_tokens, self.bound_token_id)
3
    ret.aborted_msg = self.aborted_msg
4
    return ret

它克隆出的新 matcher 共享编译产物 ctx、mask_after_eos 与 trigger/bound 配置，但不复制任何进度状态（trigger_progress、tokens_in_matcher、tokens_after_eos 全部是初值）。调度器 fork_seq 用的就是它：

1
init.structured_output.status = so_matcher.is_none() || prompt_incomplete ? 0 : 1;
2
if (!so_matcher.is_none()) {
3
    init.structured_output.matcher = so_matcher.attr("copy_without_status")();
4
}

这段在语义上完全自洽：fork 是一个「从共享前缀重新开始生成」的新序列，它的 matcher 必须回到起点，但没必要重新编译一份语法。ctx 是不可变产物，天然可以在多个 matcher 间共享；进度状态是可变的，必须各归各。

3.1 四个状态维度的分工

XGrammarMatcher 的成员变量看起来杂，其实可以分成四组，每组管一件独立的事：

最后一项 tokens_in_matcher 是回滚能力的基础：XGrammar 内部的 matcher 只记录自己的状态深度，而回滚时要决定「多少步该由内部状态机消化、多少步该由 DeepInfer 自己的簿记消化」。分开计账才可能算清楚，下面 3.5 节会看到它怎么用。

四个派生判断把状态压成布尔：

1
def _trigger_pending(self) -> bool:  return self.trigger_progress < len(self.trigger_tokens)
2
def _bound_fired(self) -> bool:     return self.tokens_after_bound >= 0
3
def constrained(self) -> bool:      return self._bound_fired() or not self._trigger_pending()
4
def finished(self) -> bool:         return self.tokens_after_eos >= 0

constrained() 是「本步要不要填掩码」的总开关，它的含义值得逐项读：

• trigger 还没匹配完、bound 也没触发 → 不受约束，模型自由生成。这是为了让模型能自然地「先说点别的，再进入格式」。强制从第 0 个 token 就进语法，在很多任务上会显著劣化生成质量——模型缺少铺垫空间，容易在 schema 的边界状态上做出生硬的跳转。
• trigger 匹配完成 → 永久受约束（trigger_progress 走到末尾后不再变化）。
• bound 触发但 trigger 未完成 → 也受约束，此时不是自由采样，而是把 trigger 的剩余 token 逐个强制（见下面 fill_vocab_mask 的第一个分支）。

3.2 触发器匹配：为什么用一个 KMP 自动机

trigger 序列要在流式生成的 token 流里被识别，而它可能出现在任意位置。最朴素的写法是「记住已经匹配了几个，下一个 token 对不上就清零」——这会漏掉重叠模式。比如 trigger 是 [A, A, B]，生成的 token 流是 A, A, A, B：朴素写法的匹配过程会在第三个 A 上清零，然后 B 上永远凑不齐，触发器失效。

DeepInfer 的做法是在构造 matcher 时预先算好 KMP 失败函数：

1
def _kmp_failure(seq: list[int]) -> list[int]:
2
    failure = [0] * len(seq)
3
    k = 0
4
    for i in range(1, len(seq)):
5
        while k and seq[i] != seq[k]:
6
            k = failure[k - 1]
7
        if seq[i] == seq[k]:
8
            k += 1
9
        failure[i] = k
10
    return failure

failure[i] 的定义是：seq[0..i] 这段前缀里，既是它的真前缀、又是它的真后缀的最长子串长度。举例，seq = [A, A, B]：

• failure[0] = 0（长度 1 的前缀没有真前缀）；
• failure[1]：[A, A] 的真前缀 [A] 也是真后缀，长度 1 → 1；
• failure[2]：[A, A, B] 的真后缀 [B]、[A, B] 都不是真前缀 → 0。

匹配时用它来「不丢进度地退让」：

1
progress = self.trigger_progress
2
while progress > 0 and token != self.trigger_tokens[progress]:
3
    progress = self._trigger_failure[progress - 1]
4
if token == self.trigger_tokens[progress]:
5
    progress += 1
6
self.trigger_progress = progress

回到 [A, A, B] 那个例子：匹配到两个 A 后 progress = 2；再来一个 A，与 B 不符，于是退让到 failure[1] = 1（保留了「一个 A 已匹配」的信息），发现 token == trigger_tokens[1] == A，于是 progress 变成 2——重叠被正确处理。这个自动机的单步复杂度是均摊 $O(1)$，整体 $O(n)$，非常适合逐 token 跑在采样路径上。

顺带说：正因为 trigger 匹配发生在生成的 token 流上，matcher 从建立到启用语法之间「不需要知道 prompt 是什么」。这也解释了 2.7 节里那条 replay_tokens 通道为什么在现有代码里没有调用方——它服务的是「生成中途恢复」的场景，而不是「prompt 里带 trigger」的场景。

3.3 掩码填充的三个分支

fill_vocab_mask 在每一步为某个请求填它那一行 bitmask，逻辑只有三个分支：

1
if self._bound_fired() and self._trigger_pending():
2
    expected = self.trigger_tokens[self.trigger_progress]
3
    vocab_mask[idx].zero_()
4
    vocab_mask[idx][expected >> 5] |= np.int32(np.uint32(1 << (expected & 31)))
5
    return True
6

7
if self.finished():
8
    vocab_mask[idx].copy_(self.mask_after_eos)
9
    return True
10

11
self.inner_matcher.fill_next_token_bitmask(vocab_mask, idx, debug_print=DEBUG_PRINT)
12
return True

分支一：bound 已触发、trigger 未走完 → 只放行一个 token。 先把整行清零（全禁），再把 expected 那一位点亮：

expected >> 5 就是 $\lfloor \text{expected} / 32 \rfloor$，expected & 31 就是 $\text{expected} \bmod 32$，两者用位运算实现除以 32 的取商取余。放行的位用 1 << bit 置到对应的 int32 字里。这样这一步的采样分布只在一个 token 上有非零概率——无论模型多想说别的，这一步只能吐出这一个 token，trigger 序列被强制逐 token 走完。

分支二：语法已收尾（EOS 已接受）→ 放行全部。 mask_after_eos 是构造时用 fill_with_full_mask=True 分配并填满 _FULL_MASK 的行，所以这一步等于「不做约束」。为什么不是「只允许 EOS」？因为语法走完之后，请求的正确行为是正常结束：模型应该输出 EOS，但如果它想多说几个字，也没有语法层面的理由阻止。真实约束交给停止条件（max_tokens、stop string）去管，语法层不再加码。这样设计避免了「语法走完后模型被逼着立刻 EOS，导致输出被截断」的经典坑。

分支三：正常状态 → 交给 XGrammar 内部状态机填。 这一支是绝大多数步走的路。注意它不需要 DeepInfer 做任何翻译工作：XGrammar 的 fill_next_token_bitmask 直接按它内部的状态表把合法位写进同一块 int32 缓冲区。

三个分支里有一个共同的约定：返回值是「有没有填成功」，而不是「允许哪些 token」；False 表示 matcher 已经终止无法再填（比如接受了 stop token 之后内部状态不允许继续），此时上层会判死请求。

3.4 mask_after_eos 只有一份，被所有请求共享

from_compiled 里分配 mask_after_eos 的动作值得单独看：

1
mask_after_eos = allocate_vocab_mask(1, vocab_size, pin_memory=False, fill_with_full_mask=True).squeeze_()

形状是 [1, ceil(V/32)]，squeeze_() 之后是一维、pin_memory=False（在 GPU 上）。它是每个 matcher 各自持有一份的，而且天然只读——所有 matcher 的这份内容完全一样（全 1）。这是「宁可多花 $\lceil V/32 \rceil \times 4$ 字节，也不引入跨请求共享状态」的取舍：一份词表 129280 的行只有约 16 KB，相对于单个请求的 KV 缓存占用可以忽略；而共享对象会带来引用计数与生命周期管理，在「请求随时可被抢占、matcher 随时可被析构」的引擎里不值得。

3.5 核心机制：投机解码下的推进与回滚

现在进入本篇的技术核心。回到 sampler._get_so_mask，看它在一个 decode 步里对每个请求做了什么：

1
if so_status[idx] == 2:
2
    for token in seq_metadata.last_accepted_token + [seq_metadata.last_output_token]:
3
        if not matcher.accept_token(token):
4
            raise RuntimeError(f"matcher rejected step token {token}, upstream invariant broken")
5
elif so_status[idx] == 1:
6
    so_status[idx] = 2
7
...
8
if matcher.constrained():
9
    mask_filled = matcher.fill_vocab_mask(self.bitmasks_cpu[0], idx)
10
    assert mask_filled, "填充 so mask 失败"
11
    so_indices_cpu[0].append(idx)
12

13
num_accepted_steps = 0
14
for k, draft_token in enumerate(seq_metadata.last_next_n_token):
15
    if matcher.accept_token(draft_token):
16
        if matcher.constrained():
17
            mask_filled = matcher.fill_vocab_mask(self.bitmasks_cpu[k + 1], idx)
18
            assert mask_filled, f"填充第 {k + 1} 步 draft so mask 失败"
19
            so_indices_cpu[k + 1].append(idx)
20
        num_accepted_steps += 1
21
    else:
22
        break
23
matcher.rollback(num_accepted_steps)

这段代码把三件不同的事按顺序叠在一起，逐个拆：

第一段：用「上一轮真正输出过的 token」把语法状态推到位。

喂进去的是 last_accepted_token + [last_output_token]，也就是上一轮实际提交给用户的全部 token：前一部分是上一轮投机验证通过的草稿 token（MtpState::accepted_tokens 的注释就写着「上一轮验证通过的 tokens（SO matcher 更新自动机时需要）」），后一部分是上一轮最后输出的那个 token。顺序不能换——语法状态机是逐 token 推进的，token 的先后就是状态的先后。

这里有一个非常关键、也很容易想错的点：被拒绝的草稿 token 绝不会出现在这个列表里。因为它在上一轮就没被提交，序列里根本没有它，语法状态自然不应该包含它。如果实现时图省事、把「这一轮看到的所有草稿」都推进状态机，那么被拒的草稿会污染状态，导致后续掩码全部错位——这是这类实现最经典的错误来源。

accept_token 返回 False 时抛出的异常信息是 upstream invariant broken（上游不变式被破坏），这个措辞精确表达了它的性质：这不是「用户输入不合法」，而是引擎内部约定被违反了——上一轮提交的 token 竟然不被语法接受，说明要么状态推进有 bug，要么验证侧的回滚没做干净。

第二段：为「本轮将要验证的草稿」预先算好每一步的掩码。

注意掩码数组的下标：so_indices_cpu[0] 对应本轮主 token 位置，so_indices_cpu[k + 1] 对应第 k 个草稿位置。这与内核侧的 step 参数是对齐的（见 3.6 节）：主 token 在 step=0，草稿在 step=k+1。

为什么草稿位置也要掩码？因为投机解码的验证逻辑在目标模型分布的层面做的：目标分布决定了某个草稿 token 的接受概率，也决定了拒绝之后从哪个残差分布重采样。如果目标分布在草稿位置上不受语法约束，那么被拒绝后的重采样就可能采样出一个语法非法的 token，约束就漏了。所以约束必须落实到每一个将被采样的位置，包括草稿位置。

而草稿位置 $k$ 的语法状态，依赖于草稿位置 $0..k-1$ 的 token 已经进入状态机。代码用 if matcher.accept_token(draft_token): num_accepted_steps += 1 else: break 表达这件事：状态机接受了这个草稿，才继续往下推；一旦某个草稿让状态机拒绝（说明它不可能是该状态下的合法后继），就停止推进——后面的草稿位置不再填掩码。这里 break 而不是 continue 是对的：状态机是因果链，第 $k$ 个草稿非法时，第 $k+1$ 个位置的状态无从谈起。

第三段：matcher.rollback(num_accepted_steps)。

这是整段代码最反直觉的一行。既然刚把草稿推进去了，为什么马上要退回来？因为草稿是投机，不是事实：

草稿 token 在验证阶段可能被接受，也可能被拒绝。语法状态机如果在这一步把它们吃进去并留在那里，那么下一轮的第一段又要把「上一轮真正输出的 token」再喂一遍——同一个 token 被消费两次，状态直接错位。

所以正确形态是「推进—填掩码—回滚」：推进只是为了在正确的状态下生成掩码（掩码是草稿位置的采样约束，必须算），回滚则是把状态机还原到「尚未见到本轮 token」的干净位置，等验证结果出来之后，下一轮用实际提交的 token 重新推进。状态机的最终一致性由下一轮的第一段保证，本轮的推进是纯粹的临时计算。

用一个具体例子把时间线钉死。设 DEEPINFER_USE_NEXTN_DECODE=5，某一步的草稿是 $[d_1, d_2, d_3, d_4, d_5]$，验证结果是前两个被接受、第三个被拒绝：

表里第 2 行与第 4 行构成一对：验证结果决定了「重新推进时喂什么」，所以回滚是安全的。回滚深度等于「状态机自己接受的草稿数」（num_accepted_steps），而它可能小于等于草稿总数——上面例子里是 3，因为 $d_3$ 虽然最终被验证拒了，但它在语法上是合法的，状态机接受了它，所以必须回滚 3 步而不是 2 步。

这里要专门强调「语法接受」与「验证接受」是两个不同的判据，两套不同的计数。前者问「这个 token 在当前语法状态下合法吗」，后者问「这个草稿 token 能否通过拒绝采样的接受判据」。$d_3$ 通过了前者但没通过后者，是完全正常的情况：语法只保证格式合法，不保证语义正确。rollback 用的是前者的计数，因为要撤销的是状态机自己推进过的那几步。

3.6 rollback 的四段式分账

rollback(steps) 的实现之所以不只有一行，是因为状态机里不止一套簿记，回滚必须按同一顺序反向拆解：

1
def rollback(self, steps: int) -> None:
2
    if self.finished():
3
        steps_after_eos = min(steps, self.tokens_after_eos)
4
        self.tokens_after_eos -= steps_after_eos
5
        steps -= steps_after_eos
6

7
    if steps == 0:
8
        return
9

10
    if self.tokens_after_eos >= 0:
11
        self.tokens_after_eos = -1
12

13
    steps_in_matcher = min(steps, self.tokens_in_matcher)
14
    if steps_in_matcher > 0:
15
        self.inner_matcher.rollback(steps_in_matcher)
16
        self.tokens_in_matcher -= steps_in_matcher
17

18
    if self.trigger_tokens:
19
        if steps > len(self._progress_history):
20
            raise RuntimeError(
21
                f"rollback {steps} steps exceeds retained trigger progress history "
22
                f"({len(self._progress_history)}, cap {MAX_ROLLBACK_TOKENS})"
23
            )
24
        self.trigger_progress = self._progress_history[-steps]
25
        del self._progress_history[-steps:]
26
    if self.tokens_after_bound >= 0:
27
        self.tokens_after_bound = max(self.tokens_after_bound - steps, -1)

四段依次是：

1. 先撤「EOS 之后」的那部分。如果当前已经处于 finished 状态，那么最后若干步都是「EOS 之后多吐出来的 token」，它们根本没进内部状态机（accept_token 在 finished() 分支里只做 tokens_after_eos += 1），所以先按 tokens_after_eos 扣减。扣完之后如果 steps 还有剩余，说明要退到「EOS 之前」——于是把 tokens_after_eos 复位为 -1（未完成）。

2. 再撤内部状态机。steps_in_matcher = min(steps, tokens_in_matcher)：只退内部状态机真正消费过的步数。这个 min 是必需的，因为内部状态机的深度可能小于总回滚步数（比如触发器的 token 还没开始进内部状态机）。

3. 再撤触发器进度。从 _progress_history 里取「第 $-\text{steps}$ 个」历史值。_progress_history 是一个只保留最近 MAX_ROLLBACK_TOKENS 项的滑动窗口：

   1
   def _record_progress(self) -> None:
   2
       if not self.trigger_tokens:
   3
           return
   4
       self._progress_history.append(self.trigger_progress)
   5
       if len(self._progress_history) > MAX_ROLLBACK_TOKENS:
   6
           del self._progress_history[0]

   为什么不重算 KMP 而要存历史？因为回滚必须 O(1)：重算意味着重新扫描整个已生成的前缀，在采样热路径上不可接受。历史窗口的上界就是回滚深度的上界，两者用同一个常量，逻辑上闭环。

4. 最后撤 bound 计数。tokens_after_bound = max(tokens_after_bound - steps, -1)：-1 是「bound 尚未触发」的哨兵值，所以减法要夹在 -1 上，回退过头也不能变成 -2。

第 3 段的 RuntimeError 是硬护栏而不是软限制：如果 steps 超过了历史窗口，说明有人试图回滚超出设计深度的步数，此时任何继续推进都会产生错误状态。把它抛出来（随后被 _get_so_mask 的 except 捕获 → set_aborted → 判死本请求），是「宁可杀一个请求，也不要产出错误结果」的直接体现。

3.7 MAX_ROLLBACK_TOKENS = nextn + 2：这个 +2 从哪来

1
MAX_ROLLBACK_TOKENS = envs.DEEPINFER_USE_NEXTN_DECODE + 2
2
...
3
self.inner_matcher = InnerMatcher(ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS)
4
...
5
self._progress_history: list[int] = []   # 上限同为 MAX_ROLLBACK_TOKENS

回滚深度必须覆盖一个 decode 步内可能推进的全部草稿，也就是 DEEPINFER_USE_NEXTN_DECODE（默认 5）。多出来的 +2 是留给边界情形的余量：一步之内除了 nextn 个草稿，还可能额外消费「EOS 之后」或「bound 边界」上的步数，而回滚的四段式分账需要历史窗口把这些都装下。这个常量同时被用在两个地方——XGrammar 内部 matcher 的 max_rollback_tokens 参数和DeepInfer 自己的 _progress_history 容量——两处必须一致，否则会出现「内部状态机还能退、但我们的簿记已经丢了」或者反过来的错配。

最后注意一个实现细节：_record_progress 里有一句 if not self.trigger_tokens: return。没有 trigger 配置时，_progress_history 永远是空列表，而 rollback 的第 3 段被 if self.trigger_tokens: 挡住——两个条件必须同进同退，否则要么白存历史、要么在不该检查的地方抛越界异常。这种「同一份开关条件在两处重复」的写法是可读性的取舍：写在一起需要一个额外的布尔成员来同步，分开写则两处各自自洽。

3.8 状态机与调度器的契约：so_status 的四个取值

结构化输出在 C++ 侧只留下两个字段：

1
struct StructuredOutputState {
2
    int status{0}; // 0 还没开始，1 刚开始（交界处），2 已经开始了，-1 aborted
3
    py::object matcher = py::none();
4
};

status 的四个取值对应四条不同的路径：

为什么要有 1 这个中间态？ 因为「最后一个 prefill chunk」是一步特殊的前向：它既消费 prompt、又产出第一个生成 token，而语法状态机在它之前从未被推进过。此时不能按 2 的路径去喂「上一轮提交的 token」（上一轮没有生成 token），也不能当作 0 完全不管（否则第一个生成 token 就不受语法约束了）。于是拆出一个明确的交界态：只填掩码，不推进状态。

交界态的判定由 llm_engine 在 prompt 步之后处理：

1
if seq_metadata.is_last_prompt_chunk and seq_metadata.sampling_params.so_args is not None:
2
    if seq_metadata.so_status == 0 and seq_metadata.so_matcher is not None:
3
        seq_metadata.so_status = 1
4
    else:
5
        logger.error("[SO] last prefill chunk with unexpected so_status %d (matcher is None: %s), killing request %s", ...)
6
        seq_metadata.so_status = -1
7
        seq.status = SequenceStatus.FINISHED_SO_FAILED

注意 else 分支：状态不符合预期就判死，而不是「尽量继续」。注释把理由写得很清楚：

1
# SO 簿记 bug 只判死本请求：直接置 FINISHED_SO_FAILED 让 process_outputs
2
# 当步回收；置 -1 仅让本步 sampler 早退（不能依赖 sampler 写回落账本，
3
# _get_so_mask 快路径可能吞掉 -1）。

这句话点出了「-1 写在哪一层」的坑：_get_so_mask 有一条快路径——当批内所有请求的 so_status <= 0 时直接返回 None，不填任何掩码。如果只把 status 置为 -1 然后指望 sampler 把 -1 写回 C++ 的账本，那么在某些步上 sampler 可能因为快路径根本没观察这个状态，-1 就永远回不去，请求会挂着一个「看起来正常」的 status 一直待着。所以终止请求这件事必须在调度器侧直接落槌（置 SequenceStatus::FINISHED_SO_FAILED），不能依赖采样侧的回报。

另一处相关的注释在 scheduler 侧的 process_outputs：

1
if (so_status_bool)
2
    seq.structured_output.status = so_status[i].cast<int>();

以及同一文件里对「判死」的说明：

1
// 判死看账本 structured_output.status 而非当步上报：sampler 对批内无活跃 SO 的步返回 None，
2
// 否则会漏掉在快路径上被吞掉的 -1

「状态写在请求的账本上，而不是当步的上报里」——这条规则值得放大来看，因为它解释了为什么 so_status 要同时存在于 Sequence::structured_output.status 与 SequenceMetadata::so_status 两个位置：前者是持久账本（跨步存在），后者是本步快照（采样器读它填掩码、回写它上报结果）。快路径意味着「本步没有 SO 请求要处理」，因此本步的上报是 None，此时任何只看「本步上报」的判断逻辑都会漏掉上一步已经写进账本的 -1。

3.9 抢占：为什么 SO 请求不能被重算

调度器的 _preempt 路径里有一条针对 SO 的硬规则：

1
// SO 请求抢占目前不支持重算，直接结束
2
const bool so_blocks_recompute = seq->structured_output.status != 0;
3
if (stop_preempted_request || so_blocks_recompute) {
4
    if (so_blocks_recompute && !stop_preempted_request) {
5
        LOG_WARN("Request {} is preempted and stops: structured output (so_status {}) can not be recomputed.",
6
            seq->request.request_id, seq->structured_output.status);
7
    }
8
    ...
9
}

抢占（preemption）是显存不足时的标准手段：把某个正在运行的请求换出，等显存够了再「重算」（recompute）它已经算过的 KV。重算的语义是「从已提交的 token 前缀重新前向一遍」——这在纯数值层面是等价的（KV 是确定性函数），但对语法状态机不成立：

• 状态机是随 step 推进的有状态对象，它的进度依赖于「当时是怎样推进的」；
• 重算路径里，采样器没有义务、也没有办法复现当时的推进序列（尤其涉及草稿推进—回滚的临时状态）；
• 更重要的是 copy_without_status() 这类接口的语义决定了状态不可回放到任意中间位置。

所以 DeepInfer 的选择是直接结束这类请求，而不是硬撑着重算。这是一个明确的工程取舍：把「SO 请求在显存压力下会被杀掉」写进设计，而不是让状态机带着可能错位的进度继续跑。日志里 can not be recomputed 的措辞也提示了它的可观测性——运维看到这条 WARN 就知道是显存压力导致的 SO 请求终止，而非用户 schema 有问题。

值得对照的是本站拆过的另两个框架的处理方式：vLLM 的 V1 引擎把「语法没编译完的请求」挡在调度之外（vLLM 源码完全拆解（十四）：结构化输出），SGLang 把「语法未就绪不进调度」做成队列纪律并配有 rollback 协同投机解码（SGLang 源码完全拆解（十）：Sampling、Grammar 与结构化输出）。DeepInfer 的差异在于：它把语法状态机当作「不可重放的有状态对象」来对待，因此在抢占这条路径上选择了保守策略。

3.10 一条不做的事：bound_token_ids 到 bound_token_id 的收敛

最后补一个接口层面的观察。so_args 的 docstring 把整条语义写成了三句话：

1
so_args: dict, "schema_type", "schema"；可选 "trigger_token_ids"（单条 list[int]）：
2
    触发 grammar 生效的 token id 序列，完整出现即触发，出现前自由生成；
3
    可选 "bound_token_id"（非负整数，0 有效；缺省或 None 表示无 bound，
4
    必须搭配非空 trigger_token_ids，列表非法）：bound 生效后不再允许自由
5
    token，先逐 token 强制走完 trigger 序列再进 grammar

把这三句话与 3.1 的 constrained() 定义对照，可以确认设计是自洽的：

• 只有 trigger_token_ids：先自由生成，trigger 完整出现后进语法；
• trigger_token_ids + bound_token_id：再加一道闸——bound token 一旦出现，就不允许再自由，必须按顺序把 trigger 余下的部分吐完再进语法。

bound_token_id 的语义是「强制通道的开关」，而不是「另一个触发器」。它的典型用法是模型输出里有一个明确的「进入结构化部分」的分隔标记（比如某个特殊 token 或固定的前缀串），一旦看到这个标记，引擎就不再信任模型的自律，直接把剩下的 trigger 部分强制走完。这也是为什么 validate_trigger_bound 要求它必须搭配非空 trigger_token_ids：

1
if bound_token_id is not None:
2
    if type(bound_token_id) is not int or bound_token_id < 0:
3
        raise ValueError("bound_token_id must be a non-negative integer")
4
    if not trigger_tokens:
5
        raise ValueError("bound_token_id requires non-empty trigger_token_ids")

没有 trigger 列表，bound_token_id 就无从定义「强制走完什么」，配置本身无意义，因此在构建前就拒掉——这属于在最早的时机做参数校验，比让它在运行期表现出奇怪行为要好得多。

4.1 它要做的事：把位图变成 $-\infty$

内核的契约由 ops.py 里的注释写死：

1
def apply_token_bitmask_by_step_inplace(
2
    logits: torch.Tensor,       # [total_q, vocab_size], fp32, packed by request, in-place output
3
    bitmask: torch.Tensor,      # [num_requests, ceil(vocab_size / 32)], int32
4
    indices: torch.Tensor,      # [num_selected_requests], CUDA int32
5
    ftp: torch.Tensor | None = None,  # [num_requests], optional CUDA int32, first packed row of each request
6
    step: int = 0,              # Offset from each request's first packed row
7
) -> None:

五个参数里，前三个是标准形态（logits、掩码、要处理的请求下标），后两个 ftp 与 step 是 DeepInfer 特有的扁平化寻址。先把前三个讲清楚：

• logits 的形状是 [total_q, vocab_size]——不是 [batch, vocab]。这是连续批处理引擎的典型布局：一个批次里每个请求的有效 query 行数不同（prompt 有 $L$ 行，decode 只有 $1$ 行），所有请求的行在显存里首尾相接成一条长表；
• bitmask 是 [num_requests, ceil(V/32)]——按请求计数，因为语法状态是请求级的（每个请求自己的状态机），所以掩码天然是「一请求一行」；
• indices 是「本步需要施加约束的是哪几个请求」的下标列表。只有真正受约束的请求才在这个列表里，这就是为什么采样器里要维护 so_indices_cpu[step] 那几个列表。

4.2 ftp 与 step：把「请求级下标」翻译成「扁平行号」

问题出在 logits 与 bitmask 的行数不一致：前者是扁平化的 $O(\sum L_i)$ 行，后者是请求级的 $N$ 行。要写对第 $j$ 个 token 的行，必须先知道「第 $r$ 个请求的第一行在扁平表里的位置」。这就是 ftp（first token position）的作用：

1
# ftp is not None 时：
2
logits_row = int(ftp[request_idx]) + step
3
logits_end = int(ftp[request_idx + 1]) if request_idx + 1 < ftp.shape[0] else logits.shape[0]
4
if logits_row >= logits_end:
5
    continue

参考实现（structure_output.py 里的 _ref 版本，用作内核的正确性基准）写得很直白：ftp[r] 是请求 $r$ 的首行行号，step 是本步在请求内部的偏移，两者相加就是本步要写的行。logits_end 用 ftp[r+1] 或表尾，用来防止越界写入——当某个请求本步没有第 step 行（比如它比别的请求短），就跳过。

这里能看出为什么掩码要按「步」组织而不是「请求」组织：投机解码下，一步之内有 $1 + \text{nextn}$ 个采样位置（主 token + 一串草稿），每个位置的语法状态不同，掩码也不同。所以采样器里准备的是 max_MTP_value + 1 份掩码数组，第 $k$ 份服务 step = k；内核用一个 step 参数在这 $k+1$ 个位置之间复用同一套寻址逻辑：

1
# 采样器侧
2
for step, indices in enumerate(so_indices):
3
    if len(indices):
4
        apply_token_bitmask_by_step_inplace(logits, so_bitmasks[step], indices, ftp, step)

4.3 内核与参考实现的对照：写对了才敢用

structure_output.py 里同时存在内核调用与一份纯 PyTorch 参考实现，两者成对出现：

1
def apply_token_bitmask_by_step_inplace_ref(logits, bitmask, indices, ftp=None, step=0) -> None:
2
    vocab_size = min(logits.shape[1], bitmask.shape[1] * _BITS_PER_BLOCK)
3
    cols = torch.arange(vocab_size, device=logits.device)
4
    word_ids = cols // _BITS_PER_BLOCK
5
    bit_ids = cols % _BITS_PER_BLOCK
6
    for request in indices.cpu():
7
        request_idx = int(request)
8
        if ftp is None:
9
            logits_row = request_idx
10
        else:
11
            logits_row = int(ftp[request_idx]) + step
12
            logits_end = int(ftp[request_idx + 1]) if request_idx + 1 < ftp.shape[0] else logits.shape[0]
13
            if logits_row >= logits_end:
14
                continue
15
        allowed = ((bitmask[request_idx, word_ids] >> bit_ids) & 1).bool()
16
        logits[logits_row, :vocab_size].masked_fill_(~allowed, -torch.inf)

逐位拆解这三行核心逻辑：

• cols // 32 与 cols % 32 把「第 $c$ 个词表 token」映射到「第几个 int32 字、字里的第几位」；
• >> bit_ids 右移到最低位，& 1 取出该位，.bool() 变成布尔张量；
• masked_fill_(~allowed, -inf) 取反后把非法位置填成 $-\infty$。用 $-\infty$ 而不是 0 或某个小负数，是因为后续要过 softmax：只有严格 $-\infty$ 的指数才是 $0$，任何有限负数都会留下非零概率（浮点上 $\exp(-100) \ne 0$），约束会被「软性泄漏」。

内核版本用 TileLang 写，通过 dispatcher 拿到：

1
kernel = ops.apply_token_bitmask_by_step_inplace_tl(vocab_size, ftp is not None)
2
kernel(logits, bitmask, indices, ftp, step)

注意 ops.py 这个 dispatcher 的设计（前面第 2 节没展开，这里补上）：它按 PLATFORM.tilelang_ops_backend 动态 import 一个后端包，然后把该包里所有以 _tl 结尾的符号平铺到自己的命名空间：

1
_pkg = importlib.import_module(f"{__package__}.{PLATFORM.tilelang_ops_backend}")
2
for _info in pkgutil.iter_modules(_pkg.__path__):
3
    _mod = importlib.import_module(f"{_pkg.__name__}.{_info.name}")
4
    for _name, _obj in vars(_mod).items():
5
        if _name.startswith("_") or isinstance(_obj, ModuleType) or not _name.endswith("_tl"):
6
            continue
7
        assert globals().get(_name, _obj) is _obj, f"backend symbol collision: {_name}"
8
        globals()[_name] = _obj

这是「接口层不做 if-else，后端差异全部收敛在包结构里」的典型写法：cuda/ 与 asc/（Ascend NPU）两个目录提供同名同签名的 xxx_tl 工厂函数，引擎代码只写 ops.apply_token_bitmask_by_step_inplace_tl(...)，完全不知道跑在哪块芯片上。那句 assert 检查符号冲突——两个模块导出同名工厂时立刻失败，避免出现「后加载的静默覆盖前一个」这种极难排查的问题。

1
class KernelUnavailableError(AttributeError):
2
    """Selected backend lacks this kernel factory; conftest.py turns it into
3
    a test skip on non-cuda backends. Subclasses AttributeError to keep the
4
    module __getattr__ / hasattr protocol intact."""

KernelUnavailableError 继承 AttributeError 并重写 __getattr__，让「某个后端没实现这个 kernel」表现为一次普通的属性缺失——这样 hasattr 探测与 pytest 的 skip 逻辑都能正常工作，而不需要每个调用点都写 try/except。这个细节说明双平台（GPU + NPU）路径是按 kernel 逐个迁移的，而不是整包切换。

4.4 正确性与性能的验证方式

配套测试（tilelang_ops/cuda/structure_output.py 的同名测试模块）做了三件事：

其一，内核 vs 参考实现的逐元素相等。 用 torch.equal 而不是 allclose：

1
logits_ref = logits.clone(); logits_tl = logits.clone()
2
apply_token_bitmask_by_step_inplace_ref(logits_ref, bitmask, indices, ftp, step)
3
apply_token_bitmask_by_step_inplace(logits_tl, bitmask, indices, ftp, step)
4
assert torch.equal(logits_ref, logits_tl)

必须是精确相等，因为掩码操作是纯位运算 + 填入 $-\infty$，没有任何浮点累加，没有理由出现数值差异。如果有差异，说明内核的位运算或寻址有 bug——用 allclose 会把这种 bug 放过去。

其二，两种寻址模式都要覆盖。 test_..._without_ftp 覆盖 ftp=None（indices 直接寻址），test_... 覆盖带 ftp 且 step ∈ {0, 2} 的情形。带 ftp 的用例特意构造了「某些请求的 qlen 恰好等于 step」以触发越界跳过分支：

1
qlen = torch.randint(step + 1, 8, (num_requests,), dtype=torch.int32, device="cuda")
2
qlen[1::3] = step   # 每三个里有一个长度正好等于 step —— 必须被跳过
3
indices = torch.tensor([0, 1, 2, 5, 8, 13], dtype=torch.int32, device="cuda")

indices 里包含 1（它属于被置为 step 的那三分之一），正好测试 logits_row >= logits_end 的 continue 分支。这是边界条件定向测试而不是随机碰撞。

其三，benchmark 记录的是带宽而不是时间。 基准用例里算的是实际 IO 字节数除以耗时：

1
io_bytes = batch_size * vocab_size * logits.element_size()
2
io_bytes += batch_size * bitmask_size * bitmask.element_size()
3
io_bytes += batch_size * indices.element_size()
4
t_us = benchmark_timer(lambda: apply_token_bitmask_by_step_inplace(logits, bitmask, indices))
5
benchmark_record(..., extras={"bandwidth_gbs": io_bytes / 1e9 / (t_us / 1e6)})

理由很直接：这个内核是纯访存型的——读一遍 logits、写一遍 logits、读一遍掩码，没有算术强度。词表 129280、fp32 时一行 logits 就是 517 KB，几百个请求就是几百 MB 的读写，耗时几乎完全由 HBM 带宽决定。所以论文/工程上关心的指标是「达到了带宽峰值的百分之几」，而不是抽象的时间数字。基准用 vocab_size=129280——这正是 DeepSeek 系模型的词表规模，说明基准是按真实模型口径设的。

顺带补一个数字直觉：掩码本身只有 $\lceil 129280/32 \rceil = 4040$ 个 int32，一行 16160 字节 ≈ 15.8 KB。相对于一行 logits 的 517 KB，掩码的带宽成本约占 3%。这就是 bitmask 方案相对于「布尔掩码」或「加法掩码」的价值：如果用布尔掩码（每 token 1 字节）就是 129 KB，用 fp32 加法掩码就是 517 KB——后者等于把 logits 的读写量翻倍。关于这一点，本站 SGLang 源码完全拆解（十）：Sampling、Grammar 与结构化输出 里有一组同样的带宽对比，可以对照阅读。

5.1 现象：同一个 prompt，batch 不同，结果不同

把话题从「离散合法性」切到「数值可复现性」。

推理服务有一个反直觉的行为：同一个 prompt、同一组采样参数、同一个模型权重，如果它被放进不同的批里跑，输出可能不一样。哪怕采样用贪心（temperature=0，每一步取 argmax）、完全无随机性，批次组成本身也会改变结果。

这不是 bug，而是 GPU 数值计算的必然结果。原因有两层，下面这张来自 Thinking Machines 的图解把第二层画得很清楚：

（图源：Thinking Machines Lab 博客 Defeating Nondeterminism in LLM Inference，https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/）

第一层：浮点加法不满足结合律。

浮点数只有有限尾数位，每次加法都要把结果舍入到可表示的值上，因此加法是一个有损、且顺序相关的运算。经典例子：$a = 1.0$，$b = 1\times 10^{-16}$，$c = -1.0$。按 $(a+b)+c$ 算，$1.0 + 10^{-16}$ 在双精度下舍入回 $1.0$，得到 $0$；按 $a + (b+c)$ 算，$b + c = -0.9999999999999999$，再加 $a$ 得到 $10^{-16}$。同一组数的和可以差出 $10^{-16}$ 量级。

单次加法误差极小，但在模型里它会被逐层放大：一个 0.001 的 logits 差异可能不改变 argmax，但一个恰好跨过决策边界的差异会直接改变采样出的 token，之后整条序列就分岔了。这也解释了为什么 batch-invariant 只在低温度 / 贪心解码 + 特定位置上表现得特别敏感——温度低的时候决策边界窄，微小差异更容易翻转结果。

第二层：并发执行的归约顺序不确定。

在 GPU 上一个矩阵乘被拆成很多线程块，每个块算一部分部分和，最后需要一次归约（reduction）把它们加起来。归约的顺序取决于哪些块恰好并发执行、以什么顺序完成，而这取决于当前的 SM 占用情况——也就取决于同批还在跑别的什么。

下面这张图把问题画在请求层面：同一个用户的请求，在确定性路径下输出只由它自己决定；在非确定性路径下，批内「其他用户的请求」也会参与进来，并改变它的结果。

（图源：Thinking Machines Lab 博客 Defeating Nondeterminism in LLM Inference）

Thinking Machines 的博客给出了一个非常实用的判据：他们发现同一批次规模下重复运行结果是稳定的，跨批次规模才不稳定。这直接把问题定位到「归约被 batch 大小影响」，而非「GPU 普遍不可复现」。基于这个观察，他们提出的修法是把归约拆成两个层次：

也就是 blog 里说的 split-K 式做法（matmul-01/02/03 系列插图对应的三种方案）：小批次路径不依赖 split-K（避免跨块归约），更大批次按固定切分累加，使每一行/每一列的结果只由它自己的数据决定。

（图源：Thinking Machines Lab 博客 Defeating Nondeterminism in LLM Inference）

5.2 batch-invariant 的准确含义

到这里可以给 batch-invariant 下一个精确的定义了：

batch-invariant 指的是：某个请求的输出，只由「它自己的输入、采样参数、模型权重」决定，与「同一个批次里还有谁、批有多大、它在批里排第几」无关。

下面这张图是「非确定性 → 确定性」的形态对照：

（图源：Thinking Machines Lab 博客 Defeating Nondeterminism in LLM Inference）

它和「确定性」不是同一件事，必须分清：

两者的关系是：batch-invariant 只保证「值不变」，不保证「无随机」。一个开启 batch-invariant 但用 temperature=1.0 的请求，它的结果仍然依赖随机数——只不过这个随机数不再由「引擎内部的计数器/时钟」提供，而由请求自己的种子提供。这就是为什么 DeepInfer 的 README 把「确定性采样」写成两条件的叠加：

Deterministic sampling: produce identical outputs for identical input tokens and sampling parameters. Set DEEPINFER_BATCH_INVARIANT=1 DEEPINFER_ENABLE_MTP_VARLEN=0 and pass a fixed seq_seed to SamplingParams (NVIDIA only).

三个条件各自负责一块：

• DEEPINFER_BATCH_INVARIANT=1：让 GPU kernel 的归约与批次无关（数值层面）；
• DEEPINFER_ENABLE_MTP_VARLEN=0：关掉变长投机（调度层面）。因为变长验证每步接受几个草稿由运行时估计器决定，而估计器读的是「实测耗时中位数」（见 第二篇），这本身就把「系统当时的负载」引进了计算路径——同样输入在不同负载下会走不同的验证深度，序列自然不同；
• 固定 seq_seed：让采样这一步的随机由请求的种子决定（随机数层面）。

5.3 seq_seed：把随机数从「引擎状态」搬到「请求属性」

sampling_params.py 里 seq_seed 的语义写在一行里：

1
seq_seed: int(uint64), the random seed of this sequence for sampling.

1
def _make_seed(seed: int | None) -> int:
2
    if seed is None:
3
        ...
4
    if isinstance(seed, int):
5
        return seed & 0xFFFF_FFFF_FFFF_FFFF

两个细节：

• 种子是 uint64，所以 & 0xFFFF_FFFF_FFFF_FFFF（即 $\bmod 2^{64}$）做规范化。校验也在：

  1
  if self.seq_seed < 0 or self.seq_seed >= 2**64:
  2
      raise ValueError(f"seq_seed must be in range of uint64_t, got {self.seq_seed}.")

• 种子是「序列级」而不是「请求级」。字段名叫 seq_seed 而不是 request_seed，因为它要跟着序列走——同一个序列在 P/D 分离的两台机器上、在重算时、在恢复时都必须是同一个种子。

种子的使用落在采样内核里。random_sampling 的签名直接暴露了这个设计：

1
def random_sampling(probs, seq_seeds, sample_positions, is_nextn_predict=False) -> torch.Tensor:
2
    """Sample one column per row, weighted by probs with each row summing to at most 1.
3

4
    Deterministic in (seq_seeds, sample_positions, is_nextn_predict): independent
5
    of batch size and row order; zero-probability columns are never sampled.
6
    """

「Deterministic in (seq_seeds, sample_positions, is_nextn_predict)」——这句话是整件事的定义。随机源不是全局的 CUDA RNG，而是由三元组 $(s, p, n)$ 通过哈希导出的：

参考实现把这一步写成了纯 PyTorch（注释里说明这是为了 Ascend 路径将来替换）：

1
def _hash_to_uniform(seed: torch.Tensor) -> torch.Tensor:
2
    """SplitMix64 hash of seed -> uniform f64 in [0, 1); int64 ops emulate uint64."""
3
    z = seed.to(torch.int64)
4
    z = z + i64(0x9E3779B97F4A7C15)
5
    z = (z ^ urshift(z, 30)) * i64(0xBF58476D1CE4E5B9)
6
    z = (z ^ urshift(z, 27)) * i64(0x94D049BB133111EB)
7
    z = z ^ urshift(z, 31)
8
    return urshift(z, 11).to(torch.float64) * (1.0 / (1 << 53))

SplitMix64 是标准的多轮位混合哈希：加黄金比例常数 $0x9E3779B97F4A7C15$，然后三轮「异或右移 + 乘大奇数」。urshift(z, 11) * 2^{-53} 取高 53 位映射到 $[0, 1)$——53 是 float64 的尾数位数，这样得到的均匀数与 float64 精度正好匹配。注意它是用 int64 运算模拟 uint64（urshift 与 i64 两个辅助函数就是干这个的），因为 PyTorch 没有原生 uint64 算术。

而三元组的混合方式是：

1
sample_case = 1 if is_nextn_predict else 0
2
seq = (sample_positions.to(torch.int64) * 3 + sample_case) * i64(0x9E3779B97F4A7C15)
3
combined = seq_seeds.to(torch.int64) + (seq ^ (seq * (1 << 13)))

三层结构值得说明：

• sample_positions * 3 + sample_case——把「位置」与「用途」编码进同一个整数。乘以 3 是因为 sample_case 有两个取值（0、1），再留一个位置；sample_case 区分「普通采样」与「NextN 草稿预测」，让同一个位置的两种用途用不同的随机数（否则草稿与主采样会共享随机性，破坏独立性）。
• * 0x9E3779B97F4A7C15 与 ^ (seq * (1 << 13))——扩散。前者是黄金比例常数乘法（把相邻的位置打散到整个 64 位空间），后者是经典的 xorshift 式扰动（x ^ (x << 13) 是 Marsaglia xorshift 的第一步）。
• seq_seeds + ...——最后把序列种子加进来，让不同序列的同一位置得到不同的随机数。

这样构造出的随机数有一个关键性质：它只依赖 $(s, p, n)$，不依赖「批里有几行、这一行排第几、谁先跑」。所以打乱批次顺序后，每行拿到的随机数不变，采样结果自然不变。

5.4 验证：三种不变性必须同时成立

random_sampling 的测试直接把这三种不变性写成了断言：

1
ref = random_sampling(probs, seq_seeds.to(torch.uint64), positions, is_nextn_predict=False)
2

3
again = random_sampling(probs.clone(), seq_seeds.to(torch.uint64), positions.clone(), is_nextn_predict=False)
4
assert torch.equal(again, ref)                      # 不变性一：重复运行相同
5

6
perm = torch.randperm(n, device="cuda")
7
out_perm = random_sampling(probs[perm].contiguous(), seq_seeds[perm].to(torch.uint64), positions[perm].contiguous(), ...)
8
assert torch.equal(out_perm, ref[perm])             # 不变性二：打乱行顺序，结果跟着打乱
9

10
for i in range(n):
11
    single = random_sampling(probs[i:i+1].contiguous(), seq_seeds[i:i+1].to(torch.uint64), positions[i:i+1].contiguous(), ...)
12
    assert single.item() == ref[i].item(), f"row {i}: single={single.item()} != batched={ref[i].item()}"
13
                                                    # 不变性三：单行跑 == 批内跑

• 不变性一（幂等）：同样的输入跑两次结果相同。它排除「内核内部用了未初始化的状态」这类问题——注意测试时 probs.clone()，强制重新读一遍显存，避免读缓存蒙对。
• 不变性二（置换等变）：把行打乱，输出也要跟着同样地打乱。这条排除了「第 $i$ 行的结果依赖它前面有几行」——正是 batch-size 依赖的本质。
• 不变性三（单行等价）：这一条最关键，因为它就是业务语义——「把请求单独跑」和「把它塞进大批里跑」必须给出同一个 token。测试对每一行都单独跑一次，逐行比对。

三条一起成立，才能说这个采样内核是 batch-invariant 的。注意测试覆盖了 h ∈ {1024, 129280}、n ∈ {256, 512, 4096}，即同时覆盖小词表和真实词表、小批和较大批。

5.5 引擎侧：不止采样器，注意力也要改

采样只是一处。日志层面上「同一输入必然给出同一输出」还要求前向计算也是批不变的，所以注意力路径也要处理。attention_backend.py 的 decode_mla_mqa_sparse 签名里有这么一项：

1
batch_invariant: bool = False,
2
...
3
output, lse = flash_mla_with_kvcache(
4
    ...
5
    enable_batch_invariant=batch_invariant,
6
)

而调用方 attn.py 直接把它接到环境变量上：

1
context, _ = decode_mla_mqa_sparse(
2
    q_b,
3
    ...
4
    batch_invariant=envs.DEEPINFER_BATCH_INVARIANT,
5
    fused_norm_rope_cast_kwargs=fused_kwargs,
6
)

enable_batch_invariant 是 FlashMLA 侧的开关，用来选择「按 batch 无关的方式组织归约」的实现路径。注意它是逐调用点显式传递的，而不是一个全局开关——这符合前面看到的平台分层风格：开关的语义（要不要 batch 不变）由引擎决定，实现（怎么做到）由内核库决定。

同样地，注意力还有一处放弃融合 kernel 的取舍：

1
self.use_fused_norm_rope_attn_rope_cast = (
2
    # disabled in dense packed-KV decoding
3
    ((not envs.DEEPINFER_BATCH_INVARIANT and cfg.engine_type == EngineType.PROMPT) or self.is_sparse or not self.has_packed_kv)
4
    and self.num_heads in (64, 128)
5
    and PLATFORM.is_gpu
6
)

这是一条布尔公式，逐个子条件读：

• not DEEPINFER_BATCH_INVARIANT and engine_type == PROMPT：开了 batch-invariant 就禁用这个融合 kernel（注释说 dense packed-KV decode 下也禁用）；
• or self.is_sparse：稀疏注意力路径可以用（稀疏路径的 token 集合由 top-k 索引显式给定，本身与批次无关）；
• or not self.has_packed_kv：没有 packed KV 时可以用；
• and num_heads in (64, 128)：融合 kernel 只支持这两种 head 数；
• and PLATFORM.is_gpu：NPU 上没这个 kernel。

为什么开 batch-invariant 要放弃融合 kernel？ 因为融合意味着「把 norm + rope + attention + rope + cast 串在同一个 kernel 里」，中间结果不落显存、不经过标准路径。融合 kernel 内部通常有自己的并行组织方式与归约顺序，且为了提高吞吐会按 batch 切分工作——这与「归约必须与 batch 无关」的目标直接冲突。放弃融合换来的是「走标准、可验证、顺序固定的实现路径」，代价是性能。这条取舍是整个 batch-invariant 路线的缩影：它不是一个免费的功能开关，而是一份用性能换可复现性的合同。

5.6 数值层的确定性：use_deterministic_algorithms

再往下还有一层。worker/runtime_config.py 里有一段进程级的内核库配置：

1
if envs.DEEPINFER_BATCH_INVARIANT or PLATFORM.is_npu:  # TODO(NPU): 防止 CI 数值不稳定
2
    deep_gemm.use_deterministic_algorithms(True)
3
    tile_kernels.use_deterministic_algorithms(True)

deep_gemm（DeepSeek 的 GEMM 库）与 tile_kernels（TileLang 内核库）都提供了 use_deterministic_algorithms 开关。开启之后，这些库会避开那些结果依赖调度顺序的实现（比如非确定性的 atomic 累加、依赖并发的部分和合并），改走顺序固定的算法。这与 5.1 节里 Thinking Machines 提出的思路是同一件事，只是落在库而不是模型代码里。

同一文件里还有几条相关约束：

1
if PLATFORM.is_npu:
2
    # TODO(NPU): 待完整支持 batch invariant 后移除
3
    assert not envs.DEEPINFER_BATCH_INVARIANT
4
if PLATFORM.is_gpu:
5
    deep_gemm.set_pdl(True)
6
    tile_kernels.set_pdl(True)

• NPU 上直接 assert 禁止 batch-invariant。这个写法很硬：不是「忽略环境变量」，而是让非法配置在启动时就崩。理由是「待完整支持 batch invariant 后移除」——因为 Ascend 路径尚未完成迁移，静默忽略会让用户以为自己的请求是确定性的，实际上不是，这比启动失败危险得多。
• GPU 上开 set_pdl(True)（Programmatic Dependent Launch，程序化依赖启动）：让后续 kernel 在前一个 kernel 还没完全结束时就可以开始做准备工作（如加载常量、初始化），从而隐藏启动延迟。PDL 与 batch-invariant 并存这件事值得注意——它说明「确定性」并不等于「放弃所有现代 GPU 优化」，被排除的是「结果依赖并发顺序」的那一类优化（如非确定性归约），而不是所有重叠手段。启动延迟优化不改变数值结果，所以可以保留。

最后，userbuffer.py 里还有一条与专家并行相关的约束：

1
allow_multiple_reduction=not (envs.DEEPINFER_BATCH_INVARIANT and envs.DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL > 0),

allow_multiple_reduction 是通信/归约路径上的一个选项：允许多次归约（通常是「先部分归约、再合并」的两段式，性能更好但顺序不固定）。batch-invariant 与「专家动态重排（EPLB）开着」同时成立时，必须关掉它——因为 EPLB 会改变专家的物理位置，如果归约顺序又随位置变化，同一 token 在不同时刻就会得到不同结果。这条约束把 batch-invariant 从「单算子正确性」抬到了「整个分布式执行路径的时序一致性」。

5.7 端到端验证：把「批内」和「单条」跑成一对

内核级的测试证明了算子不变，但真正要证明的是端到端输出一致。tests/models/test_batch_invariant.py 干的就是这件事，思路非常干净：

1
outputs = llm.generate(None, sampling_params, prompt_token_ids, use_tqdm=False)   # 64 条一起跑
2
...
3
for token_ids, output in zip(prompt_token_ids, outputs, strict=True):
4
    sp = copy.copy(sampling_params)
5
    sp.golden_token_ids = [x.token_id for x in output]      # 把批跑的结果当成「参考答案」
6
    output_single = llm.generate(None, sp, [token_ids], use_tqdm=False)[0]   # 单独跑这一条
7
    result_batched = _get_result(output)
8
    result_single = _get_result(output_single)
9
    assert result_batched == result_single

三个设计点：

第一，用真实规模配置跑。 NUM_SEQS = 64，prompt 长度按正态分布 $N(2500, 1000)$ 随机生成（max(1, ...) 保证非零），max_num_batched_tokens=16384，ep_size=8（八路专家并行），并且显式打开 DEEPINFER_USE_NEXTN_DECODE=5 让投机解码参与。投机解码必须纳入覆盖，因为它引入了额外的采样位置（草稿位置）与额外的归约路径。

第二，用 golden_token_ids 把批跑结果固定为期望。 这是 DeepInfer 一个很有意思的测试设施：SamplingParams.golden_token_ids 让这次生成沿着给定的 token 序列走（而不是自由采样）。于是「单条跑」是在验证「给定同样的决策，数值路径是否给出同样的 logprobs/top-k」。测试比对的是四元组：

1
def _get_result(output: list) -> tuple[list, list, list, list]:
2
    """Extract (token_ids, logprobs, top_token_ids, top_logprobs) from the RequestOutput list."""

比对 logprobs 与 top_logprobs 比只比对 token_ids 强得多：token id 相等可能只是「差异没跨过 argmax 边界」的巧合，而 logprob 相等是数值本身的相等。这等于把「结果一致」的判据从离散层面抬到了连续层面。

第三，ignore_eos=True 与 top_logprobs=5。 前者保证每个请求都跑满 max_tokens=50 步（不被 EOS 提前截断），让比对覆盖完整的解码长度；后者提供足够宽的比对面（每步 5 个候选的 logprob）。另外注意 prompt token 是从「非特殊 token」里随机选的：

1
special_tokens_dict = llm.llm_engine.tokenizer.added_tokens_decoder
2
non_special_tokens = [token for token in range(vocab_size) if token not in special_tokens_dict]
3
prompt_token_ids = [[random.choice(non_special_tokens) for _ in range(length)] for length in prompt_lens]

而且词表大小的取法考虑到了压缩词表：

1
vocab_size = llm.llm_engine.model_config.get_vocab_size()
2
if llm.llm_engine.model_config.use_compressed_tokenizer:
3
    from deepinfer.model_executor.models.v4_modules.engram import CompressedTokenizer
4
    vocab_size = len(CompressedTokenizer.lookup_table)

Engram 是 V4.1 里的一套压缩 token 机制（输入 token 会先映射到一张压缩表）。测试必须按压缩表的长度取随机范围，否则会生成映射表里不存在的 id——这类细节说明测试是贴着真实模型跑的，而不是在一个玩具模型上做验证。

同一目录下还有 test_ve_batch_invariant.py（多模态编码器侧的 batch-invariant），说明这条契约覆盖的不只是文本主干，还包括视觉编码器的前向。

5.8 为什么值得付这份代价

最后回到动机。batch-invariant 与确定性采样换来的是三类可用的能力：

• 前缀缓存的可判定性。 前缀缓存（prefix caching）的语义是「相同前缀复用 KV」。如果同一段文本在不同批次下算出不同的 KV，那么「命中缓存」与「未命中缓存」就会产生不同的后续输出——缓存从「透明加速」变成「改变结果的因素」。batch-invariant 让缓存变为纯粹的加速，不改变语义。这条与 第三篇 拆的两本块账直接相关。
• 可对账与可审计。 生成结果可以被离线复算：拿同样的 prompt 与 seq_seed 重跑，必须逐位相同。这对需要留证、需要复核的服务（比如合规场景下的输出存档）是硬需求。
• 回归测试的可信性。 模型/引擎的改动如果会改变输出，测试必须能稳定地报出来。如果基线本身每次跑都不同，任何 diff 都会被噪声淹没，改动只能靠人工比对——这是数值不确定最昂贵的隐性成本。

而代价是很明确的：放弃一部分融合 kernel 与并发归约优化，换取顺序固定的实现。DeepInfer 把它做成了显式开关（DEEPINFER_BATCH_INVARIANT，且只在 NVIDIA 上可用），而不是默认行为，正是因为这份代价在纯吞吐场景下不划算。这是一种典型的引擎设计取舍：把「可复现」做成按需购买的服务等级，而不是全局强加的默认值。

6.1 掩码的拷贝流：为什么不用默认流

_get_so_mask 在填完 CPU 侧的位图之后，需要把它们搬到 GPU。这段搬运被放在一条独立的 CUDA stream 上：

1
consumer_stream = torch.cuda.current_stream()
2
copy_after = torch.cuda.Event()
3
copy_after.record(consumer_stream)
4
if self.so_copy_stream is None:
5
    self.so_copy_stream = torch.cuda.Stream()
6
with torch.cuda.stream(self.so_copy_stream):
7
    self.so_copy_stream.wait_event(copy_after)
8
    so_bitmasks = [m.cuda(non_blocking=True) for m in self.bitmasks_cpu]
9
    so_indices = [torch.tensor(indices, dtype=torch.int32, device="cpu", pin_memory=True).cuda(non_blocking=True) for indices in so_indices_cpu]
10
    so_mask_ready = torch.cuda.Event()
11
    so_mask_ready.record(self.so_copy_stream)
12

13
for tensor in [*so_bitmasks, *so_indices]:
14
    tensor.record_stream(consumer_stream)
15
return so_indices, so_bitmasks, so_status, so_mask_ready

拆开看这几个动作：

• copy_after.record(consumer_stream)：在当前（计算）流上打一个事件，表示「到此为止的计算已经排好队」；

• so_copy_stream.wait_event(copy_after)：让拷贝流等待这个事件——保证拷贝的是已经算完的 logits/状态对应的掩码，而不是提前搬走；

• non_blocking=True：H2D 拷贝与计算异步进行；

• so_mask_ready 事件在拷贝流上记录，随后在采样前被等待：

  1
  if so_mask_ready is not None:
  2
      torch.cuda.current_stream().wait_event(so_mask_ready)

为什么非要一条独立流？ 因为掩码的准备是 CPU 密集的（逐请求跑状态机、填位图），而这一步被刻意安排成与 logits 计算重叠。代码里位置摆得很清楚：

1
# Start preparing the structured-output mask while logits are computed.
2
driver_inputs = sampler_inputs.driver
3
if driver_inputs is not None:
4
    so_indices, so_bitmasks, so_status, so_mask_ready = self._get_so_mask(seq_metadata_list)
5

6
# CUDA Graph model_output retains graph-batch padding; gather trims hidden_states.
7
hidden_states = gather(hidden_states, selected_token_indices)
8
logits = get_logits(hidden_states, embedding, self.vocab_size)

_get_so_mask 在 get_logits 之前调用，而它只做了 CPU 侧的填表和发起异步拷贝；GPU 侧的 logits 计算同时进行。两者的汇合点是那个 wait_event——在真正要用掩码的时候才同步。这就是本节开头所说的「结构化输出不能成为新瓶颈」的具体落地方式：不是把掩码算得更快，而是把它藏进已有的计算时间里。

最后一句 tensor.record_stream(consumer_stream) 是必备的显存安全措施：这些张量在拷贝流上被创建，但会在计算流上被使用。告诉 PyTorch 的缓存分配器「消费者是计算流」，才能避免「拷贝流用完就释放、计算流还在读」的悬空内存问题。这类细节是异步流编程里最容易埋雷的地方。

6.2 掩码缓冲区按需扩容、跨步复用

1
if not self.bitmasks_cpu or self.bitmasks_cpu[0].shape[0] < num_sampling_token:
2
    self.bitmasks_cpu = [
3
        allocate_vocab_mask(num_sampling_token, self.vocab_size, pin_memory=True) for _ in range(self.max_MTP_value + 1)
4
    ]

三个要点：

• 掩码是 pin_memory=True 的 CPU 张量。因为状态机在 CPU 上跑、位图在 CPU 上填，填完再异步 H2D。用 pinned 内存才能让 non_blocking=True 的拷贝真正做到异步（未 pinned 的内存会退化成同步拷贝）；
• 份数是 max_MTP_value + 1：一份服务主 token（step=0），max_MTP_value 份服务草稿位置。这与 4.2 节里「按步组织掩码」完全对应；
• 只扩容、不缩容：批变小了不会释放缓冲。因为缓冲区大小由「历史上见过的最大采样请求数」决定，而分配本身很便宜（一次分配终身复用，注释里在 vLLM 那篇也提到同样策略）。

注意这里没有在每步开头重置整块缓冲为「全禁」。这与本站 vLLM 源码完全拆解（十四）：结构化输出 里描述的「默认全禁、纵深防御」策略不同——DeepInfer 的做法是只填被选中的请求行，靠 so_indices_cpu[step] 精确列出要处理哪几行，未列出的行根本不读。两种策略都对，取舍在于：vLLM 的做法对「漏填」更安全（表现为卡住而不是约束失效），DeepInfer 的做法省掉了每步一次全量填充的带宽。DeepInfer 之所以敢这样，是因为它的 so_indices_cpu 与 so_bitmasks 是一对一构造的，索引列表由同一段代码生成，漏填的风险被构造过程本身排除了。

6.3 so_status 从 GPU 侧回流的路径

前面 3.8 节提到「账本优先于上报」，这里把回流的完整路径补上。采样器返回的 so_status 列表随 SamplerOutput 上到引擎，再作为参数传给调度器的 process_outputs：

1
first_token_latencies, pushed_tokens_before_finish = self.scheduler.process_outputs(
2
    ...
3
    output.so_status,
4
)

而在 C++ 侧它被写回请求账本：

1
bool so_status_bool = !so_status_obj.is_none();
2
py::list so_status = so_status_bool ? so_status_obj.cast<py::list>() : py::list();
3
...
4
for (ssize_t i = 0; i < num_output_tokens; i++) {
5
    auto &seq = *scheduler_outputs.scheduled_seqs[i];
6
    ...
7
    seq.cache.advance_optimistically(scheduler_outputs.seqlens_q[i]);
8
    if (so_status_bool)
9
        seq.structured_output.status = so_status[i].cast<int>();