论文地址

TL; DR (Too Long; Didn’t Read): CoRun 在保证推理确定性的同时，相较于 batch-invariant 方案提升了系统吞吐量。

研究背景

即使固定采样参数和随机种子，大模型的多次推理仍有可能输出不一致的答案。这个问题对模型评估和强化学习等任务带来了不利影响。例如，有论文发现多次评估同一数据集，最好的得分与最差的得分可能相差70%。

先前工作发现不确定性产生的主要原因是推理系统的动态打包（batching）和 GPU 算子的打包依赖性（batch-dependent）。推理系统的请求数不同会导致算子的输入形状不同，从而导致输出不一致。下图以求和算子为例展示了算子的打包依赖性：算子会根据输入形状自适应地调整切分（tiling）策略。当批大小为 1 时，4 个元素被分给两个流多处理器（Streaming Mutiprocessor, SM），它们各自算出 $a+b$ 和 $c+d$，再归约一次得到 $(a+b)+(c+d)$（这个做法被称为 Split-K）；当批大小为 4 时，在 batch 维度上的并行度大，于是禁用 Split-K，SM 1 计算 $x_2$ 和 $x_3$，SM 2 计算 $x_1$ 和 $x_4$，SM 2 直接计算得到 $a+b+c+d$。我们知道计算机的浮点数加法不满足结合律，因此 $(a+b)+(c+d)\neq a+b+c+d$。随着模型的尺寸变大，不确定性越来越严重。

先前工作的解决方案就是将这些涉及到归约操作的算子全部替换为打包一致性（batch-invariant）算子，禁止它们自适应地改变切分策略。如下图所示，打包一致性算子在两种批大小下都使用同一切分策略，牺牲算子时延来保证输出的一致性。但是一个固定的切分策略无法在所有批大小输入的情况下都做到最优的性能。

论文思路

打包一致性算子会引入非常大的时延开销。如下图所示，随着归约维度的增长，打包一致性算子的时延甚至会达到标准算子的 36.2 倍。当每个算子都引入额外的时延时，实验显示系统整体的吞吐量最多会下降 74%。

我们注意到，大多数标准算子虽然不是打包一致的，但它们却是位置一致的：如果算子的输入形状相同，无论一个 token 在这个 batch 中的哪个位置，它的输出都相同。这是因为输入形状相同时，切分策略一致，算子在归约维度上的归约顺序就保持一致，从而输出是一致的。

算子	是否打包一致性？	是否位置一致性？
Mean
RMSNorm
SoftMax
MatMul
Ring AllReduce
Fixed-Tree AllReduce
FlashAttention w/ Split-KV
FlashAttention w/o Split-KV

既然保证模型输入的形状就能保证一致性，那我们可以从框架层考虑，而不使用低效的打包一致性算子。

关键创新点一：CoRun 调度系统设计

我们设计了新的调度算法使得请求在预填充阶段和解码阶段都保持固定的输入形状。

1. 单独预填充：每条请求预填充时，不与其它请求打包。这保证了预填充阶段的输入一致，此时我们甚至不需要算子的位置一致性，因为输入不变可以保证输出不变。因此这个阶段我们启用了 FlashAttention 的 Split-KV 策略。
2. 固定形状解码：在解码阶段，每个请求有一个 token 参与计算，我们将输入 token 数补全（padding）至一个预先设定好的数量（即最大并发数，是同一时刻最多服务的请求数），保证输入形状的一致性，从而保证输出的一致性。值得注意的是，我们实现固定形状解码没有修改一行代码，而是利用了已有的 CUDA graph 优化：仅设置一个可用的 CUDA graph 形状（即最大并发数），当 token 数小于该形状时，系统会自动将 token 数补全到该形状。又因为这个形状一定比所有正在解码的请求数大，所以解码阶段一定是以这个形状推理的。
3. 形状对齐采样：这一部分在 CUDA graph 捕获的范围外，所以我们需要在解码阶段做手动的补全，而预填充阶段不用处理，避免冗余的计算。

虽然在解码阶段的补全操作引入了冗余的计算，确定性推理的使用场景如模型评估和强化学习都是离线批量推理任务。这种任务允许我们始终保持请求数长时间等于最大并发数，从而最小化冗余的开销。

关键创新点二：数学归纳法证明 CoRun 满足确定性

基本条件：预填充阶段不与其它请求打包，因此输出的第一个 token 是一致的。

归纳递推：假设输出的前$N-1$个 token 一致，那么在解码第$N$个 token 时，根据输入形状相同以及算子的位置一致性假设，得到第$N$个 token 也是一致的。根据数学归纳法，同一请求的输出在不同的打包请求数下是一致的。

性能评估

主要结果

在 Hy3 模型上的实验结果显示，CoRun 在保证推理确定性的同时，相较于 batch-invariant 方案提升了 15-324% 吞吐量，平均吞吐量提升一倍。

吞吐量细节

我们将预填充和解码的两类吞吐量拆解展示在下图。三类方法都是先进行预填充阶段，再进行解码阶段。非确定性方法在最快的时间内完成了两个阶段，打包一致性方法在两个阶段上的吞吐量都很慢。CoRun 由于单独预填充，它的预填充吞吐量最低，而且在前 30 秒都没有进行解码。但是，它的解码吞吐量非常高，峰值吞吐量可以达到非确定性方法峰值吞吐量的 83%，这是 CoRun 的主要收益来源。在强化学习计算（rollout）场景下，解码阶段占很长时间，CoRun 非常适合这种场景。

价值与启发

不确定性产生的根因是算子的打包依赖性，因此使用打包一致性算子是一个直接有效的方法。但本研究发现，打包一致性是一个过强的条件，如果将这个条件放宽为位置一致性，也可以通过调度的手段实现确定性，同时提升吞吐量。

需要指出的是，CoRun 仅适用于离线批量推理任务，不适用于并发度波动的在线服务场景。好消息是，我们目前已知的需要确定性推理的场景包括模型评估和强化学习都是离线批量推理任务。