Deterministic Inference Across Tensor Parallel Sizes That Eliminates Training–Inference Mismatch
Deterministic Inference Across Tensor Parallel Sizes That Eliminates Training–Inference Mismatch
作者/机构: Ziyang Zhang, Xinheng Ding, Jiayi Yuan, Rixin Liu, Huizi Mao, Jiarong Xing, Zirui Liu
A1 主要贡献
核心问题: 大语言模型(LLM)的推理过程存在不确定性,即使在贪心解码下,相同的输入在系统配置(如张量并行(TP)大小、批量大小)改变时也会产生不同的输出。这一问题源于浮点数运算的非结合性以及跨GPU归约顺序的不一致性。这种不确定性对于LLM作为评判器、多智能体系统以及强化学习(RL)等应用是致命的。特别是在RL场景中,训练引擎(通常使用TP=1的完全分片数据并行FSDP)和推理引擎(通常使用多GPU的TP以最大化吞吐量)之间的配置不匹配,会导致严重的精度差异,可能损害RL训练的性能甚至导致训练崩溃。
研究目标: 本研究旨在解决因TP规模变化导致的推理不确定性问题,实现完全确定性的LLM推理,特别是消除RL训练中训练引擎和推理引擎之间的精度不匹配问题。
创新点:
1. 识别并分析了TP导致的不一致性根源: 论文深入分析了在张量并行设置下,特别是行并行(row-parallel)层中,由于本地矩阵乘法和全局All-Reduce操作的归约顺序随TP规模变化而变化,导致了最终结果的不确定性。
2. 提出树状不变核(TBIK): 为了解决上述问题,论文提出了一套名为“树状不变核(Tree-Based Invariant Kernels, TBIK)”的TP不变矩阵乘法和归约原语。其核心思想是通过一个统一的分层二叉树结构来对齐GPU内部(intra-GPU)和GPU之间(inter-GPU)的归约顺序。
3. 实现跨框架、跨配置的位级一致性: 论文在Triton中实现了TBIK,并将其集成到vLLM和FSDP中。通过统一TBIK以及其他操作的实现(如注意力机制)和超参数(如MatMul的块大小),实现了即使在不同并行策略下,vLLM和FSDP之间也能获得位级完全相同的结果,从而为真正的“在策略(on-policy)”RL训练提供了基础。
A3 背景知识与关键挑战
2.1 IEEE 754浮点运算是非结合的
浮点运算的非结合性。一个广为人知的问题是,IEEE 754浮点运算不满足结合律,即 $(a + b) + c \neq a + (b + c)$。这意味着数字相加的顺序会因为累积的舍入误差而影响最终结果。为了缓解这个问题,采用了许多算法和系统层面的努力,例如使用融合乘加(FMA)来计算 $x*y+z$,它使用一个精确的双精度乘积,然后进行一次加法和单次舍入。在矩阵乘法(MatMul)和FlashAttention【【3】Dao, T., et al. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 2022】等核函数中,即使输入数据是较低精度的格式,通常也会使用FP32累加器来累加中间结果。其他常用操作如Softmax、LayerNorm、RMSNorm和RoPE,默认也都设置为FP32精度。
2.2 不确定性的来源
现代服务系统中的不确定性因素。在现代服务系统中,有几个因素可以改变浮点运算的顺序,从而影响最终结果。这些因素包括:(1)连续批处理(continuous batching)【【41】Yu, G.-I., et al. Orca: A distributed serving system for {TransformerBased} generative models. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 2022】,它会动态改变批次中的请求集合和批次大小;(2)操作的不同实现,例如使用Split-K与Non-Split-K矩阵乘法【【25】NVIDIA Corporation. Efficient gemm in CUTLASS. NVIDIA Developer Documentation】,由于Split-K需要累加部分和,其组合顺序因线程调度而异,从而产生不确定的结果;(3)操作的超参数,如MatMul和FlashAttention的块大小,也会改变核函数内部累加步骤的具体顺序;(4)并行系统中的集体操作,如All-Reduce,通常是不确定的,导致最终聚合的值不同;(5)并行策略,如张量并行(TP),它将工作负载分片到多个GPU上;(6)不同的GPU架构,它们可能依赖于不同的底层指令集进行矩阵乘法,例如Hopper上的wgmma。由于这些因素,先前的工作【【43】Yuan, J., et al. Give me FP32 or give me death? challenges and solutions for reproducible reasoning. CoRR, 2025a】表明,即使在贪心解码下,推理输出在不同的批次大小、GPU架构和TP配置下也可能出现显著差异。
3.1 动机:不确定性对强化学习的影响
不确定性对多场景的影响。不确定性会影响第1节中提到的多种场景。这里,我们以最流行的领域——基于LLM的强化学习(RL)为例,展示其影响。为便于说明,我们以REINFORCE算法为例,该算法通过以下方式更新策略 $\pi_{\theta}$:
其中 $a$ 是由策略网络 $\pi_{\theta}$ 生成的词元(token),$R(a)$ 是奖励。在实践中,rollout生成由vLLM和SGLang【【49】Zheng, L., et al. Sglang: Efficient execution of structured language model programs. Advances in Neural Information Processing Systems 38, 2024】等推理引擎执行,而模型训练则使用独立的后端(如FSDP)。这种混合设计按以下方式更新参数:
$\theta \leftarrow \theta + \mu \cdot \mathbb{E}_{a \sim \pi_{\text{rollout}}(\theta)} [R(a) \cdot \nabla_{\theta} \log \pi_{\text{learner}}(a, \theta)].$
在策略更新的隐式偏离。即使采样器策略 $\pi_{\text{sampler}}$ 和学习器策略 $\pi_{\text{learner}}$ 共享相同的参数 $\theta$,核函数实现和框架设置(如TP大小)的差异也会导致概率计算出现巨大差异,如图1所示。这会隐式地将一个在策略(on-policy)更新转换为离策略(off-policy)更新【【42】Yao, F., et al. Your efficient rl framework secretly brings you off-policy rl training, 2025a】。这种不匹配会破坏训练的稳定性,并对收敛和策略性能产生负面影响。值得注意的是,这个问题在专家混合(MoE)模型【【6】Fedus, W., et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 2022; 【40】Yang, A., et al. Qwen3 technical report. CoRR, 2025】中比在密集模型中更为严重,因为概率上的微小扰动可能导致路由器选择完全不同的专家。
3.2 现有工作:批量不变操作
批量不变操作(BIO)。为解决这个问题,(He & Lab, 2025)【【10】He, H. and Lab, T. M. Defeating nondeterminism in llm inference. Thinking Machines Lab: Connectionism, 2025】引入了批量不变操作(BIO),包括批量不变的FlexAttention、RMSNorm和MatMul核函数,这保证了无论批量大小如何,推理结果都保持确定性。具体来说,它通过沿着批量维度并行化计算并使每个批次元素独立计算来实现这一点,从而保证了无论批量大小如何,归约顺序都是固定的。例如,在RMSNorm中,每个批次元素都在单个计算核心上处理,消除了特征维度归约的核间通信。在MatMul中,每个核心计算一个2D块的点积,并遵循非Split-K策略在本地执行完整的归约。图3展示了其示意图。
3.3 批量不变操作与并行化的结合
BIO在不同并行策略下的表现。然而,“批量不变”只是实现完全确定性推理的第一步。如第2.2节所述,还有许多其他因素导致不确定性,其中最重要的是与并行性相关的因素。我们首先简要概述服务系统中最常用的三种并行策略,并讨论批量不变操作在每种策略下的表现。
* 数据并行(DP)。每个GPU持有一份完整的模型副本,但处理不同的批次样本子集,这实际上是沿着批量大小维度对工作负载进行分片。改变DP并行度等同于改变批量大小,批量不变操作可以确保结果保持一致。
* 流水线并行(PP)。模型的不同层分布在多个GPU或节点上,这是默认的节点间并行化策略。中间激活值按顺序传递到下一阶段以完成前向计算。改变PP度不改变每个样本内的计算顺序;因此,只要使用确定性的通信操作,模型的输出就保持确定性。
* 张量并行(TP)。每层的权重矩阵被分片到多个GPU上,层的完整输出通过聚合不同GPU上计算的部分结果获得。TP是默认的节点内并行化策略。正如我们在3.4节中分析的,改变TP配置会改变模型的输出,而批量不变核函数在处理这种变化时是无效的。我们在不同的TP大小(1/2/4/8)下运行来自AIME24数据集的相同提示,并检查输出词元序列是否保持一致。表1显示了唯一输出的平均数量,值为4表示数据集中每个提示的输出在所有TP配置下都不同。此外,如图4所示,仅由TP大小变化引起的Qwen3-8B模型输出的差异,可能导致AIME24数据集上的准确率出现超过4%的波动。
在实践中,跨环境使用不同的TP大小非常普遍,因为性能优化与TP大小是耦合的。最优TP大小通常根据所需的推理速度(吞吐量/延迟)和可用的硬件资源进行调整。TP值本身在生产环境中也可能变化以满足不断变化的性能需求。这种可变性在RL和大规模训练流水线中尤其明显,其中训练和rollout系统通常采用不同的框架和并行策略进行性能优化。如图1所示,训练通常使用TP大小为1的FSDP,而rollout引擎通常依赖于vLLM等框架,并使用较大的TP大小来最大化推理吞吐量。这种框架层面的异构性自然导致了跨环境TP配置的不匹配。
3.4 为何BIO在不同TP大小下会失效?
TP中的权重分片策略。本节我们首先解释张量并行中使用的权重分片策略,然后找出批量不变操作无法确保TP不变性的根本原因。图2展示了vLLM中Transformer模型的权重在张量并行下如何被分割。通常,TP以两种互补的方式分配工作负载:列并行和行并行。具体来说,自注意力中的QKV proj以及前馈网络(FFN)中的up proj和gate proj被实现为列并行操作。在列并行模式下(图5a),权重矩阵 $W \in \mathbb{R}^{K \times N}$ 沿着其输出维度被分割成 $C$ 个块:
$W_i \in \mathbb{R}^{K \times \frac{N}{C}}, \quad i=1, \ldots, C,$
其中 $C$ 是GPU的数量。每个GPU计算一个部分结果:
$O_i = XW_i, \quad O_i \in \mathbb{R}^{M \times \frac{N}{C}},$
最终输出通过拼接所有部分结果得到:
由于这些层扩展了特征维度,它们的输出可以保持本地分区状态,无需立即同步。
行并行操作与All-Reduce的挑战。相比之下,注意力中的o proj和FFN中的down proj是行并行的。在行并行模式下(图5b),输入矩阵 $X$ 沿着其列被分割,而权重矩阵 $W$ 沿着其行被分割,因此每个GPU持有一个切片:
$X_i \in \mathbb{R}^{M \times \frac{K}{C}}, W_i \in \mathbb{R}^{\frac{K}{C} \times N}, i = 1, \ldots, C,$
每个GPU计算部分乘积 $X_i W_i$,最终输出矩阵通过对所有GPU的部分结果求和得到:
这种设计需要一个集体all-reduce操作来合并部分结果,这引入了通信开销,并由于浮点数求和的非结合性,使得输出对求和顺序敏感。这种设计背后的原理植根于数据在Transformer层中的流动方式。列并行层(如QKV proj、up proj和gate proj)的输出作为后续行并行层(如o proj和down proj)的输入。由于列并行层自然地沿着特征维度对其输出进行分区,这些分区可以直接被行并行层消耗而无需额外同步。然而,作为自注意力和FFN模块中最后一层的行并行线性层,必须通过all-reduce操作聚合所有GPU的结果,以获得当前模块的完整输出。
TP大小改变对行并行层计算顺序的影响。改变TP大小会改变行并行层的计算顺序。对于列并行层,改变TP大小不影响输出结果,因为该层沿着 $N$ 维度分区,每个GPU产生一个大小为 $M \times \frac{N}{C}$ 的输出。这些输出是相互独立的,因为在共享维度上没有发生累加。然而,在行并行层的情况下,输入 $X$ 和权重 $W$ 都沿着 $K$ 维度在GPU之间分区。每个GPU计算一个大小为 $M \times N$ 的部分结果,行并行层通过逐元素归约聚合所有GPU的部分结果。即使归约算法本身是确定性的,参与设备的数量也会改变每个元素的计算顺序。如图7所示,当使用标准的cuBLAS GEMM进行矩阵乘法时,每个GPU首先沿着 $K$ 维度顺序执行局部归约,之后部分结果通过NCCL在同一维度上跨GPU进一步归约。因此,不同的TP大小改变了归约顺序,可能导致输出不同。因此,仅靠批量不变核函数无法保证在不同TP配置下的可复现性。
总结。为了实现不依赖于TP大小的确定性推理,有必要协同设计矩阵乘法和GPU间的All-Reduce操作,以确保行并行线性层的全局归约顺序在不同的TP配置下保持一致。
A2 方法细节
4.1 核函数设计
设计目标与核心思想。为实现TP不变的推理,我们设计了一个行并行线性层,其归约顺序与All-Reduce操作的顺序一致。现代矩阵乘法通常将输入矩阵划分为块(tile),并沿着K维度累加部分结果,如图6所示。每个大小为 $[BlockM, BlockN]$ 的输出块是通过对 $T_k = \frac{K}{BlockK}$ 个部分块求和得到的。我们的核心思想是引入一种分层的树状归约结构,该结构镜像了树状All-Reduce的累加顺序。通过为GPU内部和GPU之间的归约强制使用相同的二叉树拓扑,计算顺序在不同的TP大小下保持一致。在这种结构中,每个块对应一个叶节点,内部节点代表累加序列。归约树本身是固定的,与GPU数量无关,从而确保了跨TP配置的累加路径相同。
树状不变核(TBIK)。遵循这一原则,我们提出了树状不变核(TBIK),它对所有块应用一致的完整二叉树拓扑。每个GPU被映射到树的特定层级,GPU内部和GPU之间的归约都遵循相同的分层模式。这种设计保证了无论TP配置如何,累加过程都是确定性的,数值输出也是一致的。
4.2 实现
TBIK的两个阶段。如图7所示,我们的TBIK由两个阶段组成:使用树状MatMul核函数的GPU内归约和使用树状All-Reduce的GPU间归约。
树状归约MatMul核函数内的GPU内归约。我们使用Triton【【32】Tillet, P., et al. Triton: an intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019】实现GPU内归约。如算法1所示,每个GPU负责 $T_k/C$ 个块,其中 $C$ 表示GPU总数。为实现树状累加,必须临时存储中间部分结果。可以证明,至少需要 $L$ 个累加器,其中 $L = \log_2(T_k/C)$ 代表二叉归约树的深度。因此,我们为计算每个输出块分配一个形状为 $[L, BlockM, BlockN]$ 的累加器缓冲区 $S$。在计算过程中,来自 $A$ 的每个块 $[BlockM, BlockK]$ 和来自 $B$ 的每个块 $[BlockK, BlockN]$ 被顺序加载和相乘,部分乘积累加到 $S$ 的第一层。我们维护一个形状为 $[L]$ 的计数器张量 Count 来记录每个层级已合并的块数。每当 $l$ 层的计数器达到2时,就会触发一个进位操作:$l$ 层对应的两个部分和被归约到下一层:
$S[l+1] = S[l+1] + S[l],$
之后,$l$ 层的值和计数器都被清除。这个过程重复进行,直到处理完K维上的所有块,最终的归约输出存储在 $S[L]$ 中。对于 $T_K$ 不是2的幂的情况(例如,在Qwen3-1.7B的down proj中$K=6144$且$BlockK$是2的幂),我们引入一个自适应变量$K_{first}$,它指定在第一次进位前累加多少个块。虽然这在第一层破坏了严格的二叉归约模式,但只要 $T_K / C \times K \geq TP_{max}$,整体的不变性属性仍然成立。
使用树状All-Reduce的GPU间归约。在MatMul核函数内进行局部累加后,我们使用一个all-gather集体操作来同步每个GPU的部分结果。然后通过成对求和的方式对收集到的结果进行归约,确保GPU间的归约也遵循固定的树状结构。关于我们定制的all-reduce操作的更多细节可以在附录A.2中找到。
4.3 理论证明
操作符定义。我们通过对完整二叉树拓扑的递归来定义序列上的操作 $T(·)$:
也就是说,$T(·)$ 遵循 $T^*$ 的父子结构进行成对归约,其中每个内部节点将操作符 $\oplus$ 应用于其两个子树的输出。
定理1。设 $N$ 是总块数,$C$ 是TP大小,其中 $C$ 是2的幂。如果 $N$ 个块被均匀地划分到 $C$ 个GPU上,那么在上面定义的 $T(·)$ 操作下,无论TP大小如何,分层归约顺序是固定的。
证明概要。定义 $M = N/C$ 为每个GPU的块数,并为每个GPU $d = 1, \dots, C$ 分配连续的块:
我们证明:
$T(k_1, \dots, k_N) = T(T(\mathcal{L}_1), \dots, T(\mathcal{L}_C)),$
其中 $T(·)$ 是上面定义的二叉树操作符。
基础情况 (C=1)。当只有一个GPU时,$L_1$ 包含所有块,所以:
$T(k_1, \dots, k_N) = T(\mathcal{L}_1),$
该陈述显然成立。
归纳步骤。假设等式对 $C/2$ 个GPU成立。对于 $C$ 个GPU,将GPU分为前半部分(块 $L_1, \dots, L_{C/2}$)和后半部分(块 $L_{C/2+1}, \dots, L_C$)。根据归纳假设,每半部分的归约满足:
$T(T(\mathcal{L}_1), \ldots, T(\mathcal{L}_C)).$
因此,通过递归,分层归约执行了与全局归约完全相同的成对 $\oplus$ 操作序列,所以结果是位级一致的。□
A4 实验环境与结果
实验环境总结
- 模型: Qwen3-8B, Qwen3-32B, Mistral-7B-Instruct-v0.3, Llama-3.1-8B-Instruct。
- 数据集: AIME24 和 AMC23,用于测试数值推理和数学解题能力。
- 硬件配置: 实验在 NVIDIA RTX PRO 6000 和 NVIDIA L40S 两种GPU上进行,以验证跨硬件环境的一致性。
-
软件配置:
- 推理后端使用 vLLM。
- RL训练场景中,训练引擎使用 FSDP,推理引擎使用 vLLM。
- 所有实验在 vLLM 的 eager 模式下进行,禁用了 CUDA Graphs 和前缀缓存。
-
实验设置: 针对每个模型和数据集,在12种不同的运行时配置下进行评估,这些配置是4种TP大小(1/2/4/8)和3种批量大小(8/16/32)的组合。对于Qwen3-32B,由于内存限制,采用9种配置(TP=2/4/8, BS=8/16/32)。使用固定的随机种子和解码参数进行随机采样。
实验结果总结
实验旨在回答三个研究问题:
* RQ1. 我们的方法是否能在不同TP大小下实现位级相同的结果?
* 唯一输出数(表2): 在原生BF16推理下,结果高度不确定。仅使用BIO时,在TP=1或2时可以保证批量不变性,但当TP>2或TP大小变化时失效。而BIO+TBIK组合在所有配置下始终产生唯一的输出(平均计数为1),证明了其完全的确定性。
* 最大概率差异(表3): BIO+TBIK在所有实验设置中实现了严格为零的最大概率差异,表明模型推理达到了位级确定性。原生BF16和BIO都存在明显的概率波动。这些结果在两种不同的GPU硬件(NVIDIA L40s和RTX PRO 6000)上都得到了验证(见附录C)。
- RQ2. 我们的方法引入了多少开销?
- 核函数级性能(图8): TP-Invariant核函数比cuBLAS慢,特别是在小批量时。在大批量时,它能达到cuBLAS BF16性能的63%左右。开销主要来自树状归约所需的额外临时累加器。在FP32下,两者性能相当。
- 端到端延迟(图9): 与原生BF16相比,使用BIO+TBIK的完全确定性推理引入了显著的开销(56%到135%)。其中,BIO模块本身贡献了约18%~99%的开销,而TBIK在BIO基础上额外增加了24%~38%的开销。TBIK的开销主要源于未优化的树状All-Reduce操作(占28%~50%),而非树状MatMul核函数(仅占3%~14%)(见附录D)。
- RQ3. 我们的方法能否从根本上解决RL训练中的精度不匹配问题?
- vLLM与FSDP对齐: 作者将TBIK集成到vLLM和FSDP中,并统一了两者的线性层、注意力机制及其他核函数的实现,以消除所有不确定性来源。
- 概率差距弥合(图10): 实验测量了vLLM(TP=4)和FSDP(TP=1)生成的每个词元的概率差异。结果显示,BIO能部分减小概率差距,但仍存在明显差异。而应用TBIK后,vLLM和FSDP预测的概率变得位级相同,完全弥合了两个框架之间的差距,有效解决了精度不匹配问题。
A5 结论
本文提出了TBIK,一个用于在不同TP大小下实现确定性LLM推理的框架。通过统一的二叉树拓扑结构对齐GPU内部和GPU之间的归约顺序,TBIK消除了由浮点数非结合性和不一致的归约模式导致的不确定性。我们的方法被集成到Triton、vLLM和FSDP中,实现了跨批量大小、TP配置和框架的位级相同输出。实验证实了零概率差异和在异构硬件上的可复现性,弥合了RL流水线中vLLM和FSDP之间长期存在的精度差距。
未来的工作除了进一步优化核函数性能外,一个重要的下一步是将其保证扩展到现代量化LLM中使用的低比特位数矩阵乘法核函数。量化引入了额外的非确定性因素,如舍入、缩放和融合的反量化路径,这些都必须与TBIK中使用的相同树状累加结构对齐。设计一个能够保持确定性同时保留低比特推理效率优势的树状一致低比特GEMM,是未来工作的关键方向。
A6 附录
A.1 张量并行不变Matmul核的伪代码
算法1 张量并行不变MatMul核
要求:矩阵 A ∈ R^(M×K), B ∈ R^(K×N),块大小 BM, BN, BK,总块数 T = K/BK,第一层块数 Kfirst,归约深度 L = log2(T/Kfirst) + 1
确保:输出矩阵 C ∈ R^(M×N)
对于网格(⌈M/BM⌉, ⌈N/BN⌉)中的所有块(m, n)并行执行 do
初始化局部累加器: acc[BM][BN] ← 0
初始化临时缓冲区: S[L][BM][BN] ← 0
初始化计数器: Count[L] ← 0
for t = 0 to T - 1 do
加载块 At = A[m : m + BM, tBK : (t + 1)BK]
加载块 Bt = B[tBK : (t + 1)BK, n : n + BN]
acc ← acc + AtBt
level ← 0
while level < L do
if (level = 0 and Count[level] + 1 = Kfirst) or (level > 0 and Count[level] + 1 = 2) then
acc ← acc + S[level]
S[level] ← 0
Count[level] ← 0
level ← level + 1
else
S[level] ← acc
Count[level] ← Count[level] + 1
break
end if
end while
end for
C[m : m + BM, n : n + BN] ← acc
end for
A.2 树状All-Reduce的伪代码
自定义树状All-Reduce。我们没有直接采用NCCL内置的树状all-reduce算法,因为NCCL只允许用户指定跨节点的树拓扑,而在每个节点内部,它仍然默认使用基于链的归约。在典型的部署场景中,张量并行应用于同一节点内的多个GPU。因此,为确保节点内的归约也严格遵循树状拓扑,我们实现了一个定制的树状all-reduce算法,如算法2所示。
B 补充实验设置细节
B.1 随机采样的解码参数。根据Qwen3技术报告【【29】Team, Q. Qwen3 technical report, 2025】的官方最佳实践,推理模型的解码参数设置为(temperature=0.6, top-p=0.95, top-k=20),非推理模型的参数设置为(temperature=0.7, top-p=0.8, top-k=20)。
B.2 vLLM中的确定性推理。为在vLLM中启用确定性推理,我们所有的实验都在eager模式下进行,该模式下CUDA Graphs被禁用,前缀缓存被关闭。我们将随机种子固定为42。然而,vLLM v1引擎默认启用分块预填充(chunked prefill),并且不允许用户禁用它。这种分块策略可能与确定性推理的要求相冲突。为了高效处理长提示并在有限的GPU内存下运行,vLLM将序列的预填充过程分成多个块。这种分块执行改变了计算和调度顺序——即使底层计算核是确定性的。我们观察到,这些隐式优化确实会影响模型的实际输出。当模型大小或批量较大时,GPU内存受限。在这些条件下,推理过程并非每次运行都确定,即使在TP=1时,批量不变操作(BIO)也无法完全保证在批量大小改变时输出可复现。为了消除这种服务系统调度行为对我们方法评估的影响,我们在NVIDIA L40S上对Qwen3-32B且TP=2的实验设置了max_num_seqs=1。
C 补充可复现性评估结果
C.1 在NVIDIA RTX PRO 6000 GPU上的平均唯一答案数。如表4所示,原生BF16推理在变化的系统配置下表现出高度的不确定性。BIO可以通过减少唯一输出的平均数量来部分缓解此问题,因为在TP大小为1或2且批量大小变化时可以获得一致的结果。当将BIO与TBIK结合时,所有模型在AIME24和AMC23数据集上的平均唯一输出数降至1,表明在所有测试配置下输出完全确定。这一观察结果与我们在NVIDIA L40S GPU上的实验发现一致。
C.2 在NVIDIA RTX PRO 6000 GPU上的平均最大概率差异。表5展示了两个数据集上的平均最大概率差异。与原生BF16推理相比,BIO由于其并行策略和对归约操作的控制,略微减少了词元预测概率的波动。然而,它仍然无法处理由TP变化引起的变化。相比之下,BIO+TBIK在所有实验设置中实现了严格为零的最大概率差异,展示了位级确定性的LLM推理。这一观察结果与我们在NVIDIA L40S GPU上的实验发现一致。
D 补充性能评估结果
D.1 树状All-Reduce的开销。从图11中我们可以观察到,我们实现的树状all-reduce与NCCL all-reduce相比引入了显著的开销,范围大约在60%到180%之间。在预填充阶段影响尤其显著。这种高开销很大程度上是由于我们朴素的实现方式,可以通过在NCCL中实现定制的树状all-reduce来减少。
D.2 TBIK的细粒度延迟分解。为了更好地理解TBIK中两个核心组件——树状MatMul核和树状all-reduce操作——对整体端到端性能的贡献,我们评估了四种配置。图12展示了不同输入/输出长度下的延迟分解。虽然树状MatMul的吞吐量低于cuBLAS BF16实现,但它只引入了3-14%的开销。这种有限的影响是因为行拆分线性层只占模型总计算量的一小部分,因此其效率低下对端到端延迟的影响很小。相比之下,树状all-reduce操作贡献了更大的开销,达到28-50%。这一成本主要归因于我们当前未优化的自定义实现,以及评估环境中缺少NVLink,这显著降低了GPU间的通信效率。
E TBIK块大小选择
TBIK块大小配置。本节我们介绍了TBIK中使用的块大小,即BlockM, BlockN, 和BlockK。为实现最佳吞吐量,我们对不同数据类型采用了不同的配置,如表6所示。这些块大小是通过在NVIDIA L40S GPU上对[16, 32, 64, 128, 256]这些值进行网格搜索确定的,并为每种数据类型选择了能产生最高吞吐量的配置。
💬 评论讨论
欢迎在这里分享您的想法和见解!