Ch 09 · 批不变与确定性 — 训推一致的护身符

第二部分 · 基础设施 · 09

批不变与确定性 — 训推一致的护身符

同一个 token，不论批次多大、不论排在哪个位置、不论训练还是推理，输出 bit 完全一致。这件事看起来"显而易见"，工程上却几乎所有 GPU 加速优化都默默破坏它 —— V4 是怎么逐个 patch 回来的。

名词速通 · 一分钟看懂"批不变 + 确定性"

批不变 = 同一 token 单独跑、混在 batch 里跑、不同 batch size 跑，输出逐 bit 一致；确定性 = 同一份输入跑两次输出逐 bit 一致

一句话：大多数 GPU 优化都默认浮点结合律成立，所以会重排归约顺序换取并行度；这就破坏了"同一计算图同输入输出固定"的承诺。V4 选择放弃部分并行优化（split-K、atomicAdd），换 train / RL / infer 三段 bit-identical，让 loss spike 可调试、prompt cache 可命中、RL 优势函数可重放。

浮点结合律不成立: $(a + b) + c \ne a + (b + c)$ 在 FP32 / BF16 下都常出现 1ULP 差异。原因：每次加法做 round-to-nearest，累计误差和加法顺序绑死。这是后面所有 patch 的根源：只要顺序变了，结果就变了。
批不变（Batch Invariance）: 性质 A：同一 token 在 batch_size=1 与 batch_size=4096 下走 forward，logits 字节相同。这是 prompt cache 命中、推理 split decode 一致、训练-推理对齐的前提。
确定性（Determinism）: 性质 B：同样代码 + 同样输入 + 同样 seed，跑两遍输出字节相同。比批不变更弱（不要求 batch 变化下也一致），但更基础 —— 没有它连"重启复现"都做不到。
split-K（被放弃的优化）: 大 GEMM 把 reduction 维 $K$ 切成多块并行算，最后归约。能把利用率从 60% 提到 90%，但归约顺序由 SM 调度决定 —— batch size 一变，分块数变，顺序变，结果也变。V4 的 DeepGEMM 不用 split-K，靠其他优化追平性能。
atomicAdd（破坏确定性的元凶）: 多个 thread 并发加同一个地址，硬件保证加完不丢，但谁先加谁后加由调度决定。结果非确定。Sparse Attn 反向、MoE 反向都天然多 writer，原本最爱用它。
DeepGEMM（V4 自研 GEMM）: 替代 cuBLAS 的批不变 GEMM 实现。不用 split-K，靠 swizzle、persistent kernel、tile-level 流水把单 SM 利用率拉满。性能与 cuBLAS 持平甚至略胜。
长序列单 SM 解码 + 末尾分布式补 wave: Attention 推理的批不变方案：长 prefix 解码固定一个 SM 串行算（顺序固定），末尾 token 用 distributed shared memory 多 SM 并行补齐 —— 第二个 kernel 仍然按确定顺序归约。
SM-local 累加 + 全局确定性归约: Sparse Attn 反向的方案：每个 SM 在自己的 shared mem 里累加自己负责的部分，不跨 SM atomicAdd；最后所有 SM 各自的部分按 SM id 顺序归约。顺序固定就字节固定。
mHC 反向的 split-K（不得不用）: mHC 的 expansion 矩阵输出维只有 24，K 维 7168 → 必须 split-K。V4 的妥协：每个 split 单独 output 一个 buffer，再用确定性 reduce kernel 求和（顺序按 split id）。

一句话定位："批不变 + 确定性"是一组放弃非结合优化换可调试性 + 训推一致 + 缓存命中的工程交易。每一处都有显式代价（DeepGEMM 不用 split-K、Sparse Attn 反向不用 atomicAdd、mHC 自己写 reduce），合起来让 V4 的 1.6T 训练真正可工程化。

1. 一切问题的根源：浮点加法不结合

GPU 一切并行归约的优化都默默假设 $(a+b)+c = a+(b+c)$。但这条在 IEEE-754 浮点下从来不成立：每次加法都要 round-to-nearest，round 的方向取决于当时 mantissa 谁的位多，所以哪两个数先加会改变最后一位 bit。

📖 公式白话翻译 · 1ULP 差异是怎么来的

BF16 只有 7 bit mantissa，假设算 $1.0 + 1\mathrm{e}{-3} + 1\mathrm{e}{-3}$：

顺序 A：$(1.0 + 0.001) + 0.001$。第一步 1.001 → BF16 量化为 $1.0$（小数被吞），再加 0.001 还是 $1.0$。结果 = 1.0。
顺序 B：$1.0 + (0.001 + 0.001)$。第一步 0.002 在 BF16 下是 ~$0.001953$（精确），再加 1.0 → 1.001953 → BF16 量化 → 约 $1.00195$。

两个顺序差 ~$2\mathrm{e}{-3}$ —— 在大数 + 一堆小数的累加里这就是1 ULP 量级的偏差。看似无害，但累计 1 万次后偏差量级 $\sqrt{10^4}\times 1\text{ULP} = 100$ ULP，足以让 logits 翻不同 token，让 RL 的优势计算翻号。

所以"批不变 + 确定性"不是哲学洁癖，而是"同一段计算的语义"。要么你接受"每次跑结果不一样、loss spike 永远复现不了"，要么你按下面的方式一项一项守住顺序。

Demo · 同一组数，不同求和顺序，结果差几个 ULP

交互

读图法：四个求和都用 FP32（IEEE-754 单精度）。同一组 16 个数，A/B/C/D 仅顺序不同。"差异 ULP" = (max - min) / ULP(max)。典型差异 1–10 ULP，最坏可达 100+ ULP。
把 N 加到 64 你能看到差异变得更明显 —— GEMM 的 K 维常 4K-32K，那里的差异会更显著。这就是 split-K 一旦 K 切法变化，结果就变化的本质来源。V4 的批不变 = 把"顺序"作为一阶语义钉死。

2. 批不变：Attention 与 GEMM 各下一刀

2.1 Attention：放弃 split-KV

推理时常见的优化是split-KV：把长 prefix 的 KV 维切成多段，每段一个 SM 算 attention，最后归约。问题：batch_size 不同 → SM 占用不同 → 切法不同 → 归约顺序不同 → 同一 token 的 logits 字节不一样。

V4 选择双 kernel 方案：

长 prefix kernel：固定一个 SM 串行跑完整 prefix 的 attention。慢但顺序确定。
末尾 wave kernel：剩下的几十个 token（远小于 SM 数）用 distributed shared memory 跨 SM 协同，但归约阶段按 SM id 排序，仍然确定。

代价是 prefix 部分单 SM 跑慢约 1.5×，但因为 prefix 通常 cache 命中，绝对开销可接受；末尾 wave 阶段并行度足够，综合下来推理 throughput 比 split-KV 慢 ~10%，换来 batch-invariant logits。

2.2 GEMM：DeepGEMM 不用 split-K

cuBLAS 在大 K（≥ 4K）时默认 split-K，每个 batch 大小自动选不同切法。V4 自研 DeepGEMM 把 split-K 整体砍掉，靠下面三件事追平性能：

Persistent kernel：一个 kernel 内手动循环 N×M tile，避免反复 launch；
Tile-level pipeline：load 下一 tile 与 compute 当前 tile 重叠，让 SM 利用率不靠 split-K 也能 ≥ 85%；
Swizzle layout：shared memory 按 banks 错位排，避免 bank conflict。

实测在 V4 的典型 GEMM shape 下 DeepGEMM 比 cuBLAS-split-K 快 2–5%，并且 batch_size 任意变化字节相同。

3. 确定性：sparse attn 反向 / MoE 反向 / mHC 反向

3.1 Sparse Attn 反向：拒绝 atomicAdd

Sparse attention（CSA / HCA / Indexer）反向时，每个查询 token 的梯度要回写到它选中的那 $k$ 个 key 上。不同查询可能选中同一个 key → 多 writer → atomicAdd 是"自然"选择。但 atomicAdd 的顺序由调度决定，破坏确定性。

V4 的方案：SM-local 累加 + 全局确定性归约：

把 query token 按 SM id 划分；
每个 SM 在自己的 shared memory 里累加自己负责的查询对 keys 的贡献；
所有 SM 完成后，按 SM id 顺序把各 SM 的部分加到 global key buffer。

代价是显存占用增加（每个 SM 一份 key buffer 切片），但归约顺序由 SM id 决定，与调度无关，字节固定。

3.2 MoE 反向：token 排序预处理 + buffer 隔离

MoE 反向时，每个 expert 的梯度要回写给所有派发到它的 tokens。直接写会因为token 在 dispatch 时被打乱而非确定。V4 的做法：

每个 rank 内：dispatch 阶段对落到该 rank 的 tokens 按 (expert_id, original_token_id) 排序，得到稳定顺序；
每个 expert：单独的 grad buffer，写入位置由排序后下标决定；
合并阶段：按排序顺序遍历 buffer 累加，无 atomicAdd。

3.3 mHC 反向：被迫 split-K，但自己写 reduce

mHC 的 expansion 矩阵 $W \in \mathbb{R}^{n \times r}$ 输出维 $r=24$，K 维 $n=7168$。这个 shape 极扁，不切 K 的话单 SM 算太久 GPU 闲。所以 mHC 不得不 split-K。妥协：

每个 K-split 单独 output 到独立 buffer（不直接归约）；
归约由专门的 deterministic reduce kernel 完成，按 split id 顺序累加；
split id 与 batch size 无关，由 K 维分块决定 → batch-invariant 仍然成立。

数值演练 · 这些 patch 的代价

Attention 单 SM prefix：单 layer 推理慢 ~1.5×，但因 prefix cache 命中率高，整体 wall-time 慢 ~10%；
DeepGEMM vs cuBLAS：实测持平或快 2–5%（大 K 下持平、小 K 下快）；
Sparse Attn 反向 buffer：每 SM 一份 key 切片显存 ~40 MB（H100 SMs=132），可接受；
MoE 反向排序：dispatch 阶段额外 ~5 µs，相对 all-to-all ~480 µs 可忽略；
mHC 自定 reduce：额外 ~3% mHC 层时间，相对总训练 0.2%；

合计代价 ~3–8% wall-time。换回的是"loss spike 可逐 bit 重放" + "prompt cache 全命中" + "RL rollout 与训练对齐"，工程上是赚的。

4. 为什么这件事是隐藏 MVP

三件事少一样都做不成

调试：1.6T 模型在第 N 步爆 NaN，没有 bit-identity → 重启复现已偏；有 bit-identity → 保存 checkpoint，重新跑到那一步，attach debugger，逐 head 逐 token 看到底是哪个数先溢出。这是 V4 1.6T 工程化的根本姿态。
训推一致：训练时 logits=L_train，推理同 prompt 给 L_infer。如果两者 bit 不同，RL 的 advantage = log(π_infer / π_train) 永远不为 0，整个 rollout 数据有偏。批不变让 L_train = L_infer 字节相同，RL 的 reward 信号干净。
缓存命中：prompt cache 按 logits / KV 哈希命中。批不变让同一段 prompt 在不同 batch_size、不同部署位置命中同一份 cache，节省 30%+ prefill 算力。

所以批不变 + 确定性不是为了"理论纯洁"，而是为了让前面所有奇技淫巧（CSA / HCA / Muon / FP4 QAT / MegaMoE）可被复现、可被调试、可被部署。它是看不见的护城河 —— 平时谁都不会感谢它，出 bug 的那一天它救命。

本章小结

浮点不结合是一切非确定性的根源；GPU 一切并行归约优化默认它结合。
批不变 = 同 token 跨 batch_size 字节一致；确定性 = 同输入跨次运行字节一致。两者都不是免费的。
V4 在 Attention（双 kernel）、GEMM（DeepGEMM 不 split-K）、Sparse Attn 反向（SM-local + 顺序 reduce）、MoE 反向（排序 + buffer）、mHC 反向（split-K + 自定 reduce）逐项守住。
合计代价 3–8% wall-time，换 loss spike 可调、训推 logits 一致、prompt cache 命中。