Batch-invariant 推理完全拆解：五个推理框架如何让同一请求不再随 batch 漂移

RMSNorm：拒绝 split reduction 的诱惑

RMSNorm 的定义是

其中 $x \in \mathbb{R}^{d}$ 是一行 hidden state，$w$ 是逐通道的缩放权重，$\epsilon$ 防止除零。对一个形状为 $[B, d]$ 的张量，归约只发生在最后一维，$B$ 行之间完全独立。

最自然的并行策略是数据并行：一个 CTA 负责一行，行内用线程或 warp 做树形归约。batch 变大时，只是 grid 变大，每一行内部的归约树不变；batch 变小时，问题来了：如果 $B$ 小于 SM 数量，大量 CTA 空闲，kernel 利用率很低。一个追求性能的 kernel 工程师会自然地引入 split reduction：把一行的 $d$ 个元素切成几段，交给多个 CTA 分别求平方和，最后再合并。这样一来，归约树就依赖“切了几段”，而切几段又取决于 batch 大小，batch invariance 立刻被破坏。

batch-invariant 版本的做法很“笨”，但正确：每行固定由一个 program 处理，block 大小固定，沿列方向串行循环。vLLM 的 rms_norm_batch_invariant 就是这么写的，核心循环可以简化成下面的形式：

1
@triton.jit
2
def _rms_norm_kernel(input_ptr, weight_ptr, output_ptr,
3
                     input_row_stride, output_row_stride,
4
                     n_cols, eps,
5
                     BLOCK_SIZE: tl.constexpr, HAS_WEIGHT: tl.constexpr):
6
    row_idx = tl.program_id(0)
7
    row_start = input_ptr + row_idx * input_row_stride
8
    sum_sq = tl.zeros([1], dtype=tl.float32)
9
    # 固定 BLOCK_SIZE，沿列方向串行累加
10
    for col_offset in range(0, n_cols, BLOCK_SIZE):
11
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
12
        mask = col_idx < n_cols
13
        vals = tl.load(row_start + col_idx, mask=mask, other=0.0).to(tl.float32)
14
        sum_sq += tl.sum(tl.where(mask, vals * vals, 0.0))
15
    inv_rms = 1.0 / tl.sqrt(sum_sq / n_cols + eps)
16
    for col_offset in range(0, n_cols, BLOCK_SIZE):
17
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
18
        mask = col_idx < n_cols
19
        vals = tl.load(row_start + col_idx, mask=mask, other=0.0).to(tl.float32)
20
        out = vals * inv_rms
21
        if HAS_WEIGHT:
22
            w = tl.load(weight_ptr + col_idx, mask=mask, other=1.0).to(tl.float32)
23
            out = out * w
24
        tl.store(output_ptr + row_idx * output_row_stride + col_idx,
25
                 out.to(vals.dtype), mask=mask)

启动方式是 grid = (n_rows,)，每个 program 只处理一行。BLOCK_SIZE 被固定为 1024，循环次数只由 n_cols 决定，而 n_cols 在模型结构确定后是常数。因此无论 n_rows 是 1 还是 512，某一行的平方和都按同样的块顺序、同样的树形归约完成，结果逐位一致。代价是当 batch 很小时，GPU 上只有少量 program 在跑，算力严重闲置；对 decode 阶段动辄几十毫秒的端到端延迟来说，这点 RMSNorm 的浪费通常还能接受，因为它只占前向 FLOPs 的很小一部分。

把非批次不变的内核放进完整推理系统后，情况就反转了：服务器收到的请求集合不变时仍然确定，但单个用户无法控制同一时刻还有谁在请求，batch 组成因此成为用户侧不可见、不可控的变量。图片来自 Defeating Nondeterminism in LLM Inference。

vLLM 在实现里还处理了两个边界条件。第一，带 residual 的融合路径：rms_norm_batch_invariant 在传入 residual 时会直接调用自定义的 ops.fused_add_rms_norm，这条路径是否严格 batch-invariant 取决于该 CUDA 算子内部的归约实现，属于需要单独验证的例外，而不是“替换 aten 就自动安全”。第二，FP32 累加：kernel 内部把 bf16/fp16 输入转成 fp32 做平方和，避免 bf16 累加在 hidden 维较大时溢出，同时保证不同 batch 下的中间精度一致。

SGLang 的实现与 vLLM 同源，rms_norm 同样是 grid=(n_rows,)、BLOCK_SIZE=1024、每行一个 program；它的 batch_invariant_ops 目录还额外提供了 NPU 版本的 npu_add_rms_norm_batch_invariant，说明这套思路并不局限于 CUDA。

GEMM：把 tile 配置和 split-K 一起冻住

矩阵乘是推理里 FLOPs 最大的算子。设 $C = A B$，$A \in \mathbb{R}^{M \times K}$，$B \in \mathbb{R}^{K \times N}$，每个输出元素是

归约沿 $K$ 维。为了喂饱 tensor core，kernel 不会按单个元素计算，而是把输出切成二维 tile，例如 $128 \times 128$，每个 CTA 负责一个 tile，tile 内再用 warp 和 MMA 指令完成更细的划分。

在 batch 较大、$M$ 和 $N$ 都足够大时，数据并行就足够了：tile 数量远超 SM 数量，每个输出元素的所有 $K$ 次乘加都在同一个 CTA 内完成，归约顺序与 batch 无关。问题是 decode 阶段 $M$ 很小，可能只有几条请求、几十个 token；如果 $N$ 也不大，tile 数量不足以填满 GPU。这时 kernel 会沿 $K$ 维切分，也就是 split-K：多个 CTA 各算一部分点积，最后把部分和加起来。

图中每个输出 tile 的 K 维点积被拆成两段，分别在不同核心上累加，最后再合并。这样并行度从 4 个 tile 变成 8 个核心的负载，但每个输出元素的部分和分组取决于 split 数，split 数又随 batch 形状变化，同一行结果因此随 batch 漂移。图片来自 Defeating Nondeterminism in LLM Inference。

split-K 的问题不只是“多了一次合并”。不同的 split 数会产生不同的部分和分组，例如把 4096 个 K 元素分成 4 段和分成 8 段，合并树的形状完全不同；即使都用树形归约，舍入结果也不一样。更隐蔽的是 stream-K：它按 tile 的实际工作量动态分配 K 区间，同一个输出 tile 在不同 SM 上可能拿到不同的 K 切分，连“元素在 batch 中的位置不变”都无法保证。

要让 GEMM 满足 batch invariance，需要同时满足三个条件：

1. 每个输出元素的 $K$ 维归约在同一个 program 内完成，不跨 CTA；
2. $K$ 方向的块大小 BLOCK_K 固定，循环边界只由 $K$ 决定；
3. 不因为 $M$ 太小而切换 tensor core 指令或退回非 tensor core 路径。

注意，输出 tile 的 $M$、$N$ 形状可以随 batch 变化，只要每个输出元素内部的 $K$ 循环顺序不变。这一点容易被误解：有人以为必须把所有 shape 都固定成同一个 tile 配置，实际上不必。vLLM 的 resolve_tuned_matmul_configs() 在确定性模式初始化时，按模型的每个线性层 $(N, K)$ 提前解析好 BLOCK_SIZE_K 与 block_n 等配置，再按 $M$ 分桶选择 block_m/num_warps。同一个 $(N,K)$ 下，无论 $M$ 落在哪个桶，BLOCK_SIZE_K 都不变，每个输出元素的 K 循环次数和顺序也就固定；变的只是并行粒度和 SM 占用。

具体实现分两条路线：

• 在 SM8x（Ampere、Ada 等 8.x 设备）、ROCm 以及部分非 CUDA 后端上，cuBLAS/cuBLASLt 无法保证 batch invariance，vLLM 直接用 Triton 写的 persistent matmul 覆盖 aten::mm、aten::addmm、aten::matmul、aten::linear。这个 kernel 用固定 BLOCK_K 串行遍历 K，累加器留在寄存器里，不产生 atomicAdd，也不做 split-K。
• 在 Hopper（SM90）和 Blackwell（SM100）上，cuBLASLt 的主体实现本身是 run-to-run 确定的，batch 方差主要来自 split-K。vLLM 的做法不是替换整个 matmul，而是通过 workspace 配置关掉 split-K，并禁止 FP32 走 TF32 的近似舍入。

后者的取舍很值得体会：保留厂商库的 tile 调度和 tensor core 指令选择，只把“会随形状改变归约树”的 split-K 关掉，性能损失通常小于自己写一套 Triton kernel。前者的取舍则相反：牺牲一部分峰值性能，换取跨架构的一致性。文章给出的实测是，未经优化的 Triton batch-invariant matmul 比 cuBLAS 慢约 20%，而且在小 batch 区域由于指令过大、并行度不足，性能曲线会出现明显的“锯齿”。

除了单卡 GEMM，张量并行还会引入第四类归约：AllReduce。Ring AllReduce 按 rank 顺序累加，树形 AllReduce 按二叉树合并，两者的加法顺序不同；NCCL 还可能根据消息大小切换协议和算法。vLLM 在开启 batch invariance 时，会把自定义 AllReduce 固定为 1-stage 算法，并把 NCCL 协议钉成 Simple、通道数钉成 1、AllReduce 算法钉成 tree。原因很直接：只有 1-stage kernel 的归约顺序与数据大小、rank 数无关；多 stage 或 ring 的切分随规模变化，又会把形状依赖带回来。同样地，LoRA 的 shrink 归约被拆成 8 个固定长度的 K 分块并按固定顺序求和，避免使用 atomicAdd；在 Ada、Hopper、Blackwell 上，vLLM 还会用可打断的 CUDA Graph 让 tile 配置跟随运行时 batch，而不是被整图捕获时的形状永久冻结。

Attention：fixed split-size 与 KV 边界块

Attention 比前两者更难，因为它有两个归约维度：特征维和序列维。标准写法是

FlashAttention 把 $K$、$V$ 切成块，边算边维护 online softmax 的 running max 与 running sum。并行策略是沿 $Q$ 维切分：每个 CTA 负责一段 query，独立遍历整个 KV 序列。这个策略天然是数据并行的，因为每条 query 的归约都在自己的 CTA 内完成。

到了 decode 阶段，query 长度通常只有 1（或投机解码的几条草稿 token），沿 Q 维几乎没有并行度；如果 batch 也小，GPU 会空转。于是有了 split-KV（也叫 FlashDecoding）：把 KV 序列切成多段，每段由一个 CTA 计算部分 attention 输出与部分 softmax 统计量，最后再合并。合并公式是

其中 $m_i$、$\ell_i$、$O_i$ 分别是第 $i$ 段的 running max、running sum 和部分输出。合并本身是结合的，但浮点下不同分组会得到不同结果。

这里有一个关键的设计选择：固定 split 数量，还是固定每个 split 的大小。

常见的 split-KV 实现按“需要多少并行度”决定 split 数，再把 KV 长度平均分。例如 KV 长度 1000、需要 4 路并行，就每段 250。这样做的并行度很好，但每段的长度依赖 KV 长度和当前 batch，归约分组随之变化。FlashInfer 的 balanced scheduling 就是这种思路：选择能填满所有 SM 的最大 split size，split 数随负载变化。

batch-invariant 版本反过来：固定每个 split 的大小，例如每段 256 个 KV，最后一段可以不足 256。KV 长度 1000 时分成 256、256、256、232 四段；KV 长度 2000 时分成八段。split 数量会变，但每个 KV 位置落在哪一段、段内按什么顺序归约，只由位置决定，与 batch 无关。合并时按 split 下标从小到大做固定顺序的树形归约，就能保证同一条 query 在任何 batch 下得到相同结果。

固定 split 大小后，KV 序列的切分边界只由绝对位置决定：1000 个 KV 位置切成三段 256 加一段 232，batch 与 query 长度只改变并行任务数量，不改变每条 query 的归约树。图片来自 Defeating Nondeterminism in LLM Inference。

第二个坑是 KV cache 与当前 token 的边界块。假设 attention kernel 的 KV block size 是 32，cache 里已有 80 个 token，本次又要处理 48 个新 token。如果把 cache 和新增 token 分开计算，cache 部分需要 3 个块（两个满块加一个 16 的尾块），新增部分需要 2 个块（一个满块加一个 16 的尾块），总共 5 个块；但如果这 128 个 token 一次性 prefill，只需要 4 个满块。块数和尾块位置不同，online softmax 的合并顺序就不同。同一段 prompt 走 prefill 还是走 decode，结果会不一致。

把已缓存的 80 个 token 和新增的 48 个 token 分开算，会得到 5 个块；一次性处理 128 个 token 只有 4 个块。尾块位置不同，online softmax 的合并顺序就不同。

解决方式不是改归约公式，而是改数据布局：在进入 attention kernel 之前，先把本次要处理的 KV 写入 paged KV cache，并更新 page table（分页 KV 的块表机制见 PagedAttention 完全拆解）；这样 kernel 看到的永远是一条连续的逻辑序列，不管它是 cache 命中、chunked prefill，还是一次性 prefill，块边界都按全局位置对齐。Horace He 在文章里明确说，这一步是 attention batch invariance 的必要条件。

第三层是 采样前的 logits 归约。即使模型前向完全一致，采样器里的 log_softmax、top-k、top-p 也会对 vocab 维做归约；如果这些归约的实现随 batch 变化，最终 token 仍可能不同。vLLM 和 SGLang 都把 aten::_log_softmax、aten::softmax、aten::mean.dim 一并替换成固定顺序的 Triton 实现，原因就在这里。

把三类算子放在一起看，规律很清楚：batch invariance 不是“让浮点更精确”，而是“让每个输出元素的归约树与 batch 解耦”。RMSNorm 用固定块大小和每行一个 program；GEMM 用固定 BLOCK_K、禁用 split-K 和 stream-K；Attention 用固定 split 大小、先写 KV cache 再计算、固定 log_softmax。三者的共同点是放弃一部分形状自适应优化，换取可预测的加法顺序。

chunked prefill：切几段都要落到同一套块边界

chunked prefill 把长 prompt 切成预算大小的片段，分多次前向执行，目的是避免一个超长 prefill 长时间占住 GPU、拖高其他请求的尾延迟（调度侧的取舍可参考 Sarathi-Serve 完全拆解）。对 batch invariance 来说，它制造了一个新问题：同一条 prompt 在“一次处理”与“分三段处理”两种情形下，attention 看到的 KV 形态不同。前两段被写完、存进 KV cache；第三段只处理尾巴，并且要与 cache 里的前两段一起做 attention。如果 kernel 分别处理 cache 与新增 token，就回到了上一节说的边界块问题。

所以 chunked prefill 与 batch invariance 的兼容条件与 prefill 完全一致：每次前向都能看到从位置 0 到当前位置的完整 KV，块边界按绝对位置对齐，而不是“cache 一段、新增一段”各算各的。SGLang 在确定性推理的兼容表里把 chunked prefill 与 radix cache 逐后端列出支持情况，本质上就是同一套约束：只要 attention 的块边界按绝对位置对齐，切几段都不影响结果。

前缀缓存与 RadixCache：命中与否都要算出同样的 logits

前缀缓存（vLLM 的 prefix caching、SGLang 的 RadixCache）把公共前缀的计算结果存成 KV block，后续请求命中后直接复用；RadixCache 的树结构与引用计数在 RadixAttention 完全拆解 里有完整展开。这个优化对 batch invariance 是双重考验：

1. 同一条 prompt，首次请求要全量 prefill，后续请求只算后缀。两者的 attention 输入形态不同，但最终 logits 必须相同，否则“缓存命中”会改变用户看到的答案。
2. 多租户共享前缀时，block 会被不同请求复用，block 在 cache 里的物理位置与请求的逻辑位置脱钩。如果 kernel 依赖物理块连续性或按块枚举顺序做归约，复用就会带来差异。

SGLang 的确定性测试专门有 radix_cache 模式：同一段 prompt，一次走缓存命中、一次走全量 prefill，比较输出是否逐位相同。它的确定性推理文档把 RadixCache 支持情况列在特性表里，说明这条路径是被显式验证的，而不是“理论上没问题”。

投机解码：草稿长度不能改变目标模型的 logits

投机解码先让草稿模型生成 $k$ 个候选 token，再由目标模型一次前向验证，接受最长前缀；特征级草稿与树验证的机制可参考 EAGLE 完全拆解。草稿长度 $k$ 随上下文变化（例如 EAGLE-2 的动态草稿树），验证时的 query 长度也随之变化：有时是 1 个 token 的普通 decode，有时是一次验证 6 或 8 个 token。

对 batch invariance 而言，危险点在于 query 长度也进入了 kernel 的并行维度。如果 attention kernel 按 $Q$ 维切分、且切分方式依赖于 query 总数，那么“验证 6 个 token”与“逐个 decode”在归约树上就不一样。解决办法仍然是固定归约结构：Q 维切分只影响“哪些 query 放在一起算”，但每条 query 对 KV 的遍历顺序必须只由 KV 位置决定；对被拒绝的 token，必须把 KV 回滚到接受点，避免残留的草稿 token 改变后续 attention 的可见范围。

SGLang 的确定性测试集里专门有 EAGLE 拒绝采样的用例，说明推理框架并不把“确定性”和“投机解码”当作天然兼容的两件事，而是逐条路径单独验证。

专家并行与 EPLB：AllReduce 和专家重排

MoE 模型的推理常用专家并行（EP）：不同 rank 持有不同专家，token 需要被 dispatch 到目标 rank，计算结果再 combine 回来。这一步涉及两类通信：

• AllToAll dispatch/combine：token 从各 rank 发出，按专家分桶后送到目标 rank。如果分桶顺序、每个 rank 收到的 token 排列随 batch 变化，后续的加权求和顺序就会变化。
• AllReduce：MoE 层的输出可能需要在 TP 组内做部分和合并，AllReduce 的算法（ring/tree/1-stage）决定加法顺序。

EPLB（Expert Parallelism Load Balancer）在运行中复制热点专家、重排专家位置，以均衡负载；DeepInfer 的副本分配与逐层放置可参考 MoE 专家并行与 EPLB。重排本身不改变数学结果（每个 token 只会被路由到逻辑上同一个专家），但它会改变“哪个 rank 持有哪个专家”，进而改变通信路径与合并顺序。对 batch invariance 来说，这意味着专家重排不能改变逻辑专家的归约顺序：物理副本可以换，逻辑上的加法树必须保持稳定。

一个具体的例子是 DeepEP 的 allow_multiple_reduction 参数：它允许在多次归约中采用不同策略以换取带宽，这会破坏确定性。DeepInfer 在开启 batch invariance 且启用专家重排时，把这个参数强制设为关闭，原因正是两者冲突。

调度顺序与批内位置

还有一类更细的不确定性来自批内位置。如果一个 kernel 的 tile 调度按 pid 递增分配，而 pid 与请求在 batch 中的位置相关，那么请求 A 排第 1 位和排第 7 位，可能落到不同的 tile 编号上；只要 kernel 内部的归约与 tile 编号无关，结果就不受影响。但如果 kernel 使用了按 tile 编号变化的 splitable 调度（stream-K）或按 tile 分组的 persistent 调度，批内位置就会进入归约路径。DeepInfer 的 EPLB 实现可参考 DeepInfer 源码完全拆解（四）：MoE 专家并行与 EPLB。

因此“batch-position invariance”与“batch invariance”要分开验证。多数厂商 GEMM 满足前者；Triton persistent matmul 如果采用固定的 tile 映射，也能满足前者。真正的风险点集中在 stream-K、balanced split-KV 和动态 EPLB。

采样器：种子也必须可复现

模型前向确定性之后，采样器还有自己的随机性。temperature、top-k、top-p 采样需要随机数。多数框架的随机数生成器按“batch 内位置”分配种子或偏移量，batch 变化时同一条请求拿到的随机数就不同。要做到端到端确定性，采样器必须支持按请求绑定种子：例如 SGLang 在确定性推理模式下默认使用固定采样种子 42，并允许请求携带 sampling_seed；DeepInfer 要求传入固定的 seq_seed，并关闭投机解码的变长草稿路径。

到这里可以把系统层的约束归纳成一张表：

vLLM：把 aten 算子整体替换成 batch-invariant 实现

vLLM 的入口是一个环境变量：

1
VLLM_BATCH_INVARIANT=1 vllm serve meta-llama/Llama-3.1-8B-Instruct

vLLM 的文档目前把这项能力标注为 beta，支持 NVIDIA compute capability 8.0 及以上、AMD gfx90a/gfx942/gfx950，以及需要 Triton Attention 后端的 Intel XPU；并且给出了一份已验证模型清单（DeepSeek-V3/V3.1/R1、Qwen3 系列、Llama 3 系列、GPT-OSS 等）。内部初始化流程分四步，定义在 vllm/model_executor/determinism/batch_invariant.py 的 init_batch_invariance() 里：

1. resolve_tuned_matmul_configs()：按模型维度提前选好 block_k 等参数，避免运行时再按 batch 调参；
2. override_envs_for_invariance()：覆盖 AllReduce、cuBLAS workspace、NCCL 协议、ROCm skinny GEMM、torch.compile 的 inductor batch_invariant 配置；
3. enable_batch_invariant_mode()：注册 torch.library.Library("aten", "IMPL")，用 Triton 实现覆盖 aten::mm、aten::addmm、aten::mm.out、aten::addmm.out、aten::matmul、aten::linear、aten::mean.dim、aten::_log_softmax、aten::softmax、aten::_softmax、aten::bmm、aten::bmm.out；
4. 关闭 TF32 与 fp16/bf16 的 reduced-precision reduction，防止 matmul 在累加时降精度。

这段代码里有几个设计细节很值得讲。

按硬件分流。在 Hopper 和 Blackwell 上，vLLM 不替换 matmul，而是设置 CUBLAS_WORKSPACE_CONFIG=:16:8、CUBLASLT_WORKSPACE_SIZE=1 来关闭 split-K；注释里写得很清楚：“Hopper (SM90) and Blackwell (SM100): the only source of batch variance is split-k”。在 SM8x 和 ROCm 上则退回到 Triton persistent matmul，因为 cuBLASLt 和 hipBLASLt 在这两类设备上无法保证 batch invariance。同一套 API，不同硬件走不同路径，这是工程上非常典型的分层。

注册 .out 重载。PyTorch 的 torch.mm(a, b, out=c) 走的是 aten::mm.out 而不是 aten::mm；如果只替换前者，代码里带 out= 的路径仍然会落到厂商库。vLLM 把 .out 变体一起注册，并在 torch 2.12+ 上对内置 bmm 使用 allow_override=True，说明“把所有入口封死”是这件事的必要步骤。

对 XPU 的差异化处理。注释指出 oneDNN 后端的 softmax、log_softmax、bmm 本身就 batch-invariant，所以只在 CUDA 上注册这些覆盖；XPU 上只替换 matmul/linear 相关入口与 mean。这说明“哪些算子是问题”本身依赖后端，不能一刀切。

AllReduce 的固定。override_envs_for_invariance() 把 VLLM_CUSTOM_ALLREDUCE_ALGO 固定为 1stage，并设置 NCCL_LAUNCH_MODE=GROUP、NCCL_PROTO=Simple、NCCL_MIN/MAX_NCHANNELS=1、NCCL_ALGO=ring,tree;allreduce:tree。注释里有一处非常具体的坑：NCCL 2.31 之后，设置 NCCL_ALGO 会清空所有集合通信的算法表，只保留被点名的算法；如果集合通信没有被点名，会以 ncclInvalidUsage 失败。因此代码先把 ring、tree 为所有集合通信打开，再把 AllReduce 钉到 tree。这种“看起来无关的配置导致启动失败”的细节，正是确定性工程比论文复杂的地方。

测试体系。vLLM 的 tests/v1/determinism/ 下有一整套用例：test_batch_invariance、test_online_batch_invariance、test_rms_norm_batch_invariant、test_matmul_batch_invariant、test_nvfp4_batch_invariant、test_mx_linear_batch_invariant、test_cutlass_batch_invariance、test_lora_shrink_batch_invariant。从命名可以看出它覆盖了量化 GEMM、LoRA、在线服务等不同路径，也说明 vLLM 把“batch invariance”当作一个持续维护的子系统，而不是一个开关。

SGLang：deterministic inference 模式与采样种子

SGLang 的入口是一个启动参数：

1
python3 -m sglang.launch_server \
2
    --model-path Qwen/Qwen3-8B \
3
    --attention-backend fa3 \
4
    --enable-deterministic-inference

它的文档直接写明：实现基于 Thinking Machines Lab 的 batch_invariant_ops，并在其上兼容 chunked prefill、CUDA Graph、RadixCache 与非贪心采样。代码层面，python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py 提供了：

• matmul_kernel_persistent 与 _matmul_persistent_triton、_matmul_persistent_deepgemm 两条 matmul 路径，其中 DeepGEMM 路径会先调用 get_deterministic_algorithms() 检查状态——注释写明 DeepGEMM 0.2 的 bf16 GEMM 默认走 cuBLASLt、其归约会随 batch 变化——必要时临时切到确定性算法，算完再切回去；
• _log_softmax_kernel / log_softmax，固定 BLOCK_SIZE 沿 vocab 维串行；
• mean_kernel / mean_dim；
• bmm_kernel_persistent / bmm_batch_invariant；
• _rms_norm_kernel / rms_norm_batch_invariant，注释里标明改编自 vLLM 的对应实现；
• enable_batch_invariant_mode()，用 torch.library 覆盖 aten::_log_softmax、aten::rms_norm 等，并提供 set_batch_invariant_mode() 上下文管理器；
• get_batch_invariant_attention_block_size() 返回 block_m=16, block_n=16，用于把 attention 的块大小固定住。

与 vLLM 相比，SGLang 的特点在于：

attention 后端受限。确定性推理只在 FlashInfer、FA3、Triton 三个后端下支持，文档里的特性表写明 FlashInfer 后端不支持 RadixCache，而 FA3 和 Triton 全部支持。这里必须点出来：RadixCache 支持与否直接影响“缓存命中路径与全量 prefill 路径是否一致”，不是可以忽略的表格项。

采样种子按请求绑定。确定性推理模式下默认种子是 42，并且请求可以携带 sampling_seed。对 GRPO 这类需要“同一 prompt 多次采样、但每次可复现”的场景，这个设计很实用：不同 seed 给出多样本，同一个 seed 给出逐位相同的样本。

RadixCache 一致性验证。python3 -m sglang.test.test_deterministic --test-mode radix_cache 专门比较“缓存命中”与“未命中”的输出，预期结果是 Unique samples: 1。这个测试设计值得借鉴：它验证的不是“同一个脚本跑两遍”，而是两条不同的运行时路径。

NPU 分支。python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py 覆盖了 aten::_log_softmax、npu_fused_infer_attention_score、npu_add_rms_norm 等 NPU 算子，说明同一套思路可以在昇腾上复现，但需要逐算子重写。

TensorRT-LLM：静态形状与 CUDA Graph 的另一条路

TensorRT-LLM 在公开文档与代码组织里没有 --enable-deterministic-inference 这样的开关，但它有一套结构性的“确定性来源”：构建期把模型编译成静态形状的引擎，运行时按固定 shape 执行。

传统 TensorRT 引擎在构建时确定每个 kernel 的 tile 配置、tensor core 指令与 split 策略，运行时不能根据 batch 大小重新选 kernel。加上 CUDA Graph，整个前向是一次图的回放，kernel launch 顺序、grid 大小都固定。这天然满足了“同一 shape 重复执行结果一致”，也满足“在同一个引擎、同一个 shape 下 batch 组成不影响单条请求”的很多条件。

但要讲清楚它的边界：

• 跨 batch size 不保证一致。TensorRT-LLM 的 dynamic shape 或 padding 策略可能会把不同 batch 形状的请求映射到同一个执行 profile；如果两个 profile 使用不同的 kernel 配置，同一个请求在不同 profile 下结果就可能不同。它不是 vLLM 那种“任意 batch 组成逐位一致”的承诺。
• 插件与自定义 kernel 是变量。attention 插件、MoE 插件、量化 GEMM 插件各自有自己的归约实现；插件升级、量化方案切换都会改变数值。
• P/D 分离与 KV 传输。TensorRT-LLM 的 disaggregated serving 需要把 KV cache 从 prefill 实例搬到 decode 实例；如果两边使用不同的引擎配置，KV 的量化与布局可能不同，结果也会不同。

因此，TensorRT-LLM 的“确定性”更接近“在固定引擎、固定 shape profile 下的可重复”，而不是推理框架级的 batch invariance。它的优势是把很多不确定性提前到构建期消除；劣势是灵活性下降，动态 shape 与多 profile 会让“同一请求走哪个 profile”变成新的不确定性来源。如果要在 TensorRT-LLM 上做真正的 batch-invariant 服务，通常需要把 batch 维度固定成少量 profile，并保证所有 profile 的数值行为一致，或者把请求 padding 到固定形状——后者代价是算力浪费。

llama.cpp：单序列语义与 ubatch 的边界

llama.cpp 的目标不是高并发在线服务，而是本地、多平台的推理。它的确定性特性来自结构设计而非显式开关：

每个序列独立构图。在 llama_context 中，llama_batch 被切分成若干 llama_ubatch，每个 ubatch 构建自己的 llm_graph_result。对于单个序列，同一段 prompt 的图结构只由 token 与 KV 状态决定；如果 n_ubatch 足够大、整段 prompt 落在一个 ubatch 里，结果就是固定的。

ubatch 切分是主要变量。当 prompt 长度超过 n_ubatch（默认等于 n_batch），它会被切成多个 ubatch 分段计算。这与 chunked prefill 类似：只要 attention 在每一段都能看到前面的 KV，且 KV 的块边界按位置对齐，结果就应当稳定；但不同模型、不同 backend 的 attention kernel 是否真的如此，需要逐项确认。

多序列共享 ubatch。llama.cpp 的 ubatch 可以把多个序列的 token 合并在一起，以提升 CPU/GPU 利用率（split_simple 等拆分方式）。这时，不同序列共享同一个图，但每条序列的 attention 只访问自己的 KV；是否 batch-invariant 取决于各 backend 的 kernel 实现。llama_context.cpp 里有一条注释：“TODO: add new split mode where we pad the input sequences so that ubatch.equal_seqs == true”，说明作者也意识到“等长序列填充”可以简化 batch 的处理，只是当前的 split 方式并不保证这一点。

sampling 与 logits 的边界。llama.cpp 的采样器按序列各自维护 sampler chain，logits 的提取位置与 batch 内位置无关；但 logits 的 n_vocab 维归约由 backend 的 softmax/top-k 实现决定，不同 CPU 指令集（AVX2、AVX512、NEON）的归约可能给出不同结果。所以 llama.cpp 上的“确定性”通常是“同一机器、同一后端、同一 ubatch 配置下可重复”，跨设备/跨后端不保证逐位一致。

适合的场景。本地单用户、agent 本地运行、评测离线批跑；在这些场景里，请求队列通常很短，batch 组成稳定，ubatch 切分固定，确定性带来的风险通常小于在线服务。但如果把 llama.cpp 的 server 当作多用户服务，且 batch 动态变化，就需要自己核对 n_batch、n_ubatch、attention 后端的数值行为。

DeepInfer：batch invariant、确定性采样与 EPLB 的耦合

DeepInfer 是 DeepSeek 开源的推理引擎，仓库见 deepseek-ai/DeepInfer。当前版本支持 DeepSeek-V4.1-Flash，面向 NVIDIA SM100 与华为 Ascend 950。它的 README 把 “Batch-invariant inference” 与 “Deterministic sampling” 单列为两项功能，并给出了很具体的配置方式：

1
# batch invariance
2
export DEEPINFER_BATCH_INVARIANT=1
3

4
# deterministic sampling：需要同时满足三个条件
5
export DEEPINFER_BATCH_INVARIANT=1
6
export DEEPINFER_ENABLE_MTP_VARLEN=0
7
# 并向 SamplingParams 传入固定的 seq_seed

这两组配置揭示了 DeepInfer 对确定性的分层理解：前向的 batch invariance 与采样器的确定性是两件独立的事。前向确定性保证了 logits 不随 batch 变化；采样确定性还要求随机数种子按请求绑定（seq_seed），并且关闭“变长投机验证”（MTP_VARLEN）——因为变长验证会改变 query 长度和 KV 访问路径，README 明确要求确定性采样时把它关闭。

代码层面，deepinfer/environ.py 用 EnvVar.bool() 定义了 DEEPINFER_BATCH_INVARIANT；deepinfer/worker/runtime_config.py 的 configure_worker_kernels() 在 GPU 上调用 deep_gemm.use_deterministic_algorithms(True) 与 tile_kernels.use_deterministic_algorithms(True)，同时在 Ascend 上直接断言不允许开启（注释是“待完整支持 batch invariant 后移除”）。这说明 DeepInfer 的确定性建立在它自研的算子库上：DeepGEMM 与 TileKernels 提供确定性模式，而不是覆盖 PyTorch 的 aten。

与 EPLB 的耦合值得单独展开。deepinfer/model_executor/parallel_utils/userbuffer.py 在分配 DeepEP 的 EP 缓冲区时写：

1
allow_multiple_reduction=not (
2
    envs.DEEPINFER_BATCH_INVARIANT
3
    and envs.DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL > 0
4
)

含义是：只要同时开启了 batch invariance 与专家周期性重排，就不允许 DeepEP 在一次通信中做多种 reduction。原因在系统层小节已经讲过：多种 reduction 的合并顺序会随分组变化，专家重排又会改变分组，两者叠加就破坏了逻辑上的归约顺序。

fused kernel 的让步。deepinfer/model_executor/models/v4_modules/attn.py 里有一个 use_fused_norm_rope_attn_rope_cast 开关，启用条件包含 not envs.DEEPINFER_BATCH_INVARIANT and cfg.engine_type == EngineType.PROMPT。也就是说，在 prompt 引擎的 dense packed-KV 路径上，确定性模式会关闭默认高性能的“norm + RoPE + attention + RoPE + cast”融合 kernel，退回更朴素、归约顺序可控的实现。这是一个很典型的“确定性成本”落点：融合越深，越难保证每一步的归约顺序与 batch 无关。

attention 后端的开关透传。deepinfer/model_executor/layers/attention_backend.py 的 flash_mla_with_kvcache(..., enable_batch_invariant=batch_invariant) 把开关传给 FlashMLA；对 sparse MLA 路径，DeepInfer 还会走 TileLang 的 flash_mla_sparse_fwd 与 fused_norm_rope_attn_rope_cast.decode；后者在 fused_norm_rope_cast_kwargs 不为空时直接走融合 kernel。需要说明的是，sparse 路径的融合判断条件与 dense prompt 路径不同，确定性开关并不能自动覆盖所有融合分支，这也是“逐路径验证”比“打开一个开关”更重要的原因。

测试。tests/models/test_batch_invariant.py 的做法是：64 条随机 prompt（长度按均值 2500、标准差 1000 采样，max_tokens=50、top_logprobs=5），先以 batch 形式生成，再让单条 prompt 单独跑一遍，比较 token id、logprob、top-5 token 与 top-5 logprob 是否完全相等；单独跑时借用首轮的 token 作为 golden_token_ids，让两条路径对齐在同一条随机流上。用例在 EP 规模为 8 的集群上运行，并可开启 use_nextn_decode=5 覆盖 NextN 投机解码。这个测试比“同脚本跑两遍”严格得多：它同时验证了批内位置、batch 大小、EP 并行与投机解码对结果无影响。

把五个框架并排看，可以归纳出四种工程策略：

1. 算子替换型（vLLM、SGLang）：在 aten 或算子库层注入确定性实现，调度和模型代码基本不变。优点是通用性强、可以覆盖任意模型；缺点是依赖每个算子都被正确覆盖，融合算子与量化路径容易漏。
2. 引擎编译型（TensorRT-LLM）：在构建期冻结形状与 kernel 选择，运行时只回放固定图。优点是数值路径稳定；缺点是对动态形状不友好，跨 profile 不保证一致。
3. 端侧结构型（llama.cpp）：依靠单序列构图与固定的 ubatch 切分，确定性与平台和后端强相关。
4. 引擎自研型（DeepInfer）：从自研算子（DeepGEMM、TileKernels、FlashMLA）开始就提供确定性模式，并在调度、通信、专家重排处显式处理耦合关系。优点是全链路可控；缺点是跨平台（尤其 Ascend）仍在补齐。

自己动手验证：三条检查线

无论使用哪个框架，验证 batch invariance 至少要过三条线：

第一条：同一请求，不同 batch 组成。 构造一组短 prompt、一组长 prompt，混合成不同 batch（例如 1、4、16、64 条一起），比较同一条 prompt 的 token id、logprob、top-k logprob。DeepInfer 的 test_batch_invariant.py 就是这个思路的完整版本。

第二条：不同并行路径。 让同一条 prompt 分别走“全量 prefill”与“前缀缓存命中”，以及“普通 decode”与“投机解码验证”。SGLang 的 --test-mode radix_cache 就是这条线。

第三条：不同 chunk 切分。 同一条长 prompt 分别以整段 prefill 与 chunked prefill 计算，比较 logits。可以借用 vLLM 的 tests/v1/determinism/test_online_batch_invariance.py 的思路：通过并发数控制 batch 大小，比较同一请求在不同负载下的输出。

一个最小化的验证脚本可以写成：

1
from openai import OpenAI
2

3
client = OpenAI(api_key="EMPTY", base_url="http://localhost:8000/v1")
4

5
def ask(prompt, seed=42):
6
    r = client.completions.create(
7
        model="meta-llama/Llama-3.1-8B-Instruct",
8
        prompt=prompt, max_tokens=64, temperature=0.0, seed=seed,
9
        logprobs=5,
10
    )
11
    return r.choices[0].logprobs
12

13
ref = ask(TARGET_PROMPT)
14
# 在另一段时间、不同并发负载下再次请求，比较 logprobs 是否逐位相同

如果比较的是 logprob 数组而不是最终文本，就能在 token 分叉之前发现问题。对比文本“看起来一样”并不可靠：分歧往往发生在第 100 个 token 之后，而评测只比较最终答案。