Towards Deterministic Inference in SGLang and Reproducible RL Training

发表时间: 2025-09 · Blog post by The SGLang Team (lmsys.org)

原文: https://lmsys.org/blog/2025-09-22-sglang-deterministic

The SGLang Team

速读

一句话结论
SGLang 联合 slime 框架，通过引入批次无关算子并改造分块预填充策略，实现了完全确定性的 LLM 推理与 100% 可复现的强化学习训练，且在开启 CUDA graphs 时将性能开销压缩至仅 34.35%。

要解决什么问题
大语言模型推理中的不确定性会隐式地将同策略（on-policy）强化学习转变为异策略（off-policy）强化学习，严重破坏了严谨科学实验所需的可靠性。过去，即便用户在 SGLang 中将采样温度硬性设定为 0，由于动态批处理（dynamic batching）和基数缓存（radix cache）机制的存在，推理结果依然会发生随机波动。这个问题的核心卡点在于底层算子（如 RMSNorm、矩阵乘法、注意力机制）的归约（reduction）过程对批次大小高度敏感：当同一个请求被随机组合进不同大小的批次时，内核中归约拆分的过程会随之改变，导致每次归约块的计算顺序和大小出现差异。由于计算机底层的浮点数加法不具备严格的数学结合律，这种微小的计算顺序改变会引发浮点误差的随机累积与放大。只要批次大小在动态变化，这种底层机制就会导致非确定性的输出，成为阻碍强化学习训练完全复现的最大障碍。

怎么做的
核心思路是替换底层对批次敏感的算子，并在系统调度层面对齐计算块的边界，从而在数学层面保证浮点数计算顺序的绝对一致。该方案由四个关键部件构成，各自承担不同的职责：
第一，引入批次无关（batch-invariant）内核。系统直接集成了 Thinking Machines Lab 的均值、对数 softmax 和矩阵乘法算子，并针对多种后端实现了固定切分大小（split-KV size）的注意力内核。具体而言，在 FlashInfer 后端中利用 fixed_split_size 和 disable_kv_split 参数在内核规划阶段固定切分大小；在 FlashAttention 3 后端中将切分数量强制设为 1；在 Triton 后端中固定解码的切分大小，借此甚至将确定性推理扩展到了 AMD 硬件上。
第二，改造分块预填充（chunked prefill）策略。原有的尽力而为分块策略会为了凑齐最大块大小（如 8192），将一个序列的 KV 块随机截断，这破坏了注意力内核所需的固定切分大小。改造后的逻辑强制将截断点对齐到 split_kv_size 的整数倍：

$$ L_{truncate} = n \times split\_kv\_size $$

这保证了被推迟到下一个分块计算的序列片段依然能作为一个完整的单元被注意力内核处理，维持归约行为的确定性。
第三，实现可复现的非贪婪采样。为了在温度大于 0 时依然保持确定性，废弃了原生的 torch.multinomial，引入全新的 multinomial_with_seed 算子。它通过一个带种子的哈希函数生成 Gumbel 噪声来扰动 logits：

$$ Output = \arg\max (Logits + Gumbel(Hash(inputs, seed))) $$

使得相同的输入和种子对永远产生相同的采样结果，满足强化学习 rollout 的随机性需求。
第四，在 slime 框架中打通全链路。后向传播回退到 Flash Attention v2 以保证梯度确定性，开启 Megatron 的确定性模式，并强制设定 NCCL_ALGO=Ring、NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 等环境变量以及 PyTorch 的确定性算法开关，最终实现端到端的严格复现。

效果如何
实验在 Qwen3-8B 模型上进行，硬件使用 H100 80GB（用于吞吐量测试）和 H200 140GB（用于离线推理延迟测试）。对比基线为 SGLang 原有的非确定性常规模式（Normal Mode），分别在 FlashInfer、FlashAttention 3 和 Triton 三种注意力后端下进行了严格对比。
在确定性验证测试（包含单提示词、长短混合、同源不同前缀三种场景，共 50 次采样）中，常规模式会产生 2 到 18 种不同的输出，而确定性模式在所有后端和场景下均严格输出 1 种唯一结果。在 slime 框架的 GRPO 强化学习训练中，两次独立的训练运行产生了完全一致的 rollout 响应和首批迭代损失曲线，实现了 100% 的复现。
在性能方面，开启 CUDA graphs 后，确定性推理的吞吐量实现了约 2.8 倍的加速（例如 FlashInfer 后端从 441.73 tokens/s 提升至 1245.51 tokens/s）。在模拟 256 个请求、覆盖 1024 到 8192 不同输入输出长度的离线推理负载中，FlashInfer 和 FlashAttention 3 后端的平均性能损耗仅为 34.35%（显著优于早期方案的 61.5%）。
作者也坦承了该方法的局限性与代价：确定性模式的运行速度依然明显慢于常规模式，主要瓶颈在于批次无关的矩阵乘法和注意力内核尚未充分优化；目前仅支持稠密模型，暂不支持 DeepSeek-V3 等 MoE 模型；且在基准测试中，由于 FlashInfer 和 Triton 的基数缓存支持仍在开发中，该功能被暂时禁用；此外，大于 TP2 的张量并行设置仍需进一步改造归约内核才能保证确定性。

主要贡献

大语言模型推理结果的不确定性会隐式地将同策略（on-policy）强化学习转变为异策略（off-policy）强化学习，这为严谨的科学实验带来了挑战。为了解决这一问题，本文致力于在开源推理引擎中实现稳定且可用的确定性推理，并进一步实现完全可复现的强化学习（RL）训练。

本文的核心创新点与贡献如下：

构建完全确定性的推理引擎：基于Thinking Machines Lab（TML）的批处理不变算子（batch-invariant operators），SGLang成功实现了完全确定性的推理。
广泛的特性兼容性：该确定性推理机制能够与分块预填充（chunked prefill）、CUDA图（CUDA graphs）、基数缓存（radix cache）以及非贪婪采样（non-greedy sampling）等核心特性保持完美兼容。
显著的性能提升：借助CUDA图，SGLang实现了2.8倍的加速，并将确定性推理的性能开销从TML报告的61.5%大幅降低至仅34.35%。
实现100%可复现的RL训练：SGLang与slime团队合作，以极少的代码修改解锁了完全可复现的强化学习训练。在Qwen3-8B模型上的验证实验表明，两次独立的训练运行能够产生完全相同的损失曲线，为科学实验提供了所需的可靠性。

背景知识与设计原则

大语言模型推理一致性的重要性：大语言模型（LLMs）推理输出一致性的能力正变得越来越重要。例如，正如研究人员所指出的，推理结果的不确定性会隐式地将同策略（on-policy）强化学习（RL）转变为异策略（off-policy）强化学习。然而，在SGLang中，即使将温度（temperature）参数调至0，由于动态批处理（dynamic batching）和基数缓存（radix cache）的使用，采样过程仍然不是确定性的。

批处理大小变化导致的不确定性：正如Thinking Machines Lab博客中所提到的，不确定性的最大来源是不断变化的批处理大小（batch sizes）：即使用户反复提交相同的提示词（prompt），由于该请求可能与其他用户的请求被打包在不同的批次中，批处理大小的差异会导致非确定性的推理结果。

归约拆分过程与浮点数非结合性：为了进一步解释，不同的批处理大小会影响内核（kernels）的归约拆分（reduction splitting）过程。这导致每个归约块的顺序和大小发生变化，由于浮点算术的非结合性，从而引发非确定性的输出。为了解决这个问题，他们将归约内核（如RMSNorm、矩阵乘法、注意力机制等）替换为批处理不变（batch-invariant）的实现。这些内核也被作为一个配套库发布，供外部集成使用。

He, Horace and Thinking Machines Lab, "Defeating Nondeterminism in LLM Inference"

SGLang的确定性推理解决方案：基于Thinking Machines Lab的工作，SGLang提供了一个强大且高吞吐量的确定性LLM推理解决方案，该方案结合了批处理不变内核、CUDA图、基数缓存以及分块预填充（chunked prefill）技术，并具备高效的性能。其确定性已经通过全面的测试和RL训练实验得到了广泛验证。

SGLang的核心增强功能：

集成批处理不变内核：集成了来自Thinking Machines Lab的批处理不变内核，包括均值（mean）、对数归一化指数函数（log-softmax）和矩阵乘法内核。
实现批处理不变的注意力内核：实现了具有固定拆分键值（split-KV）大小的批处理不变注意力内核。支持多种后端，包括FlashInfer、FlashAttention 3和Triton。
全面兼容常见推理特性：与分块预填充、CUDA图、基数缓存等常见推理特性完全兼容，当启用确定性推理时，这些特性均保持支持状态。
暴露单请求种子：在采样参数中暴露了每个请求级别的种子（per-request seed），允许用户即使在温度大于0的情况下也能启用确定性推理。
更优的性能表现：与TML博客中报告的61.5%的性能下降相比，SGLang在使用FlashInfer和FlashAttention 3后端时，平均性能下降仅为34.35%，这代表了显著的改进。与最小化集成相比，结合CUDA图可以实现2.8倍的加速。

方法细节

分块预填充（Chunked Prefill）的挑战：SGLang的分块预填充技术旨在管理具有长上下文的请求。然而，其默认的分块策略违反了注意力内核对确定性的要求。

输入序列与拆分单元的示例说明：如图所示，考虑两个输入序列seq_a和seq_b，每个序列的上下文长度为6,000。分块预填充的最大块大小（chunk size）为8192，而确定性注意力所需的拆分键值大小（split-KV size）为2,048。每个序列可以被划分为三个较小的单元（a1至a3，以及b1至b3），其长度分别为2,048、2,048和1,904。如果这些较小的单元在分块预填充期间保持完整，那么它们就可以被同一个注意力内核处理，并产生确定性的归约行为。

对齐截断点以保证确定性：标准的分块策略遵循"尽力而为"（best-effort）的原则。在这个例子中，该策略试图通过将seq_b的b2单元拆分为两个较小的部分，来生成一个包含8,192个标记（tokens）的chunk_1。这会导致截断点不一致，因为拆分后b2的长度取决于seq_a的长度。为了解决这个问题，我们调整了分块逻辑，将截断点与split_kv_size的整数倍对齐。这种调整确保了对b2的处理被推迟到后续的块中，从而允许它作为一个完整的单元被注意力内核计算。

注意力后端（Attention Backends）的定制化修改：注意力内核是实现确定性的重要组成部分。针对不同的注意力后端，我们以不同的方式对其进行了修改，以满足其使用要求。

FlashInfer后端：对于FlashInfer后端，我们利用了批处理不变FA2内核中的fixed_split_size和disable_kv_split参数，在内核规划（kernel planning）期间固定拆分大小。分块预填充的截断与预填充拆分大小对齐。
FlashAttention-3后端：对于FlashAttention-3后端，将flash attention内核的num-splits固定为1，以确保确定性。
Triton后端：对于Triton后端，我们固定了解码（decoding）的拆分大小，并手动设置了分块预填充的对齐大小。借助Triton后端的全平台可扩展性，确定性推理也可以在AMD硬件上运行。

可复现的非贪婪采样（Reproducible Non-Greedy Sampling）：为了将确定性扩展到贪婪解码（greedy decoding）之外，我们引入了一个新的采样函数：multinomial_with_seed。

基于种子哈希函数的Gumbel噪声扰动：我们不再依赖在批处理下固有非确定性的torch.multinomial，而是使用由基于种子的哈希函数（seeded hash function）生成的Gumbel噪声来对逻辑值（logits）进行扰动。因此，即使在温度大于0的情况下，相同的(inputs, seed)对也始终会产生相同的采样结果。

保留强化学习所需的随机性：这一修改在实现了确定性的多项式采样（deterministic multinomial sampling）的同时，保留了强化学习部署（rollouts）所必需的随机性。

强化学习框架集成（RL Framework Integration - slime）：我们将温度大于0的确定性推理集成到了slime的GRPO训练流程中。在初步实验中，重复的RL训练运行在前几次迭代中产生了完全相同的rollout响应和损失值，证实了rollout过程本身是确定性的。

实现完整训练可复现性的关键配置：在后续的更新中，我们通过实现以下关键配置，进一步启用了完整的训练可复现性：

Flash Attention：使用Flash Attention v2而不是v3，以启用确定性的反向传播（backward passes）。
Megatron：设置--deterministic-mode标志以进行确定性训练。
环境变量（Environment Variables）：配置关键设置：
- NCCL_ALGO=Ring
- NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
- CUBLAS_WORKSPACE_CONFIG=:4096:8
PyTorch：启用torch.use_deterministic_algorithms(True, warn_only=False)。

实现真正的端到端RL训练管道：通过这些全面的更改，我们成功在slime中实现了GRPO的完全训练可复现性，从而启用了真正确定性的端到端RL训练管道。

实验环境与实验结果

实验1：验证确定性行为

实验内容：引入了一个确定性测试套件，以验证在不同批处理条件下推理结果是否保持一致。测试包含三个子测试：
1. Single（单一测试）：在不同的批处理大小下运行相同的提示词，检查输出是否相同。
2. Mixed（混合测试）：在同一批次中混合不同类型的提示词（短提示词和长提示词），并验证一致性。
3. Prefix（前缀测试）：使用源自同一长文本但前缀长度不同的提示词，随机进行批处理，测试结果在多次试验中是否可复现。
实验环境：模型使用QWen3-8B。启用了CUDA graph和分块预填充。对于Flashinfer和Triton禁用了基数缓存（因其支持仍在开发中）。
实验结果与分析：在50次采样试验中，统计了观察到的唯一输出数量（数值越低表示确定性越高，1代表完全确定）。
- 在Normal（常规）模式下，FlashInfer、FA3和Triton在Single测试中分别产生了4、3、3种不同输出；在Mixed和Prefix测试中也产生了多种不一致的输出（最高达18种）。
- 在Deterministic（确定性）模式下，无论使用哪种注意力后端，在所有子测试（Single, Mixed, Prefix）中，唯一输出数量严格保持为1。结论证实该方法完美实现了确定性行为。

实验2：CUDA Graph加速评估

实验内容：对比启用和禁用CUDA graph时，确定性推理的总吞吐量（tokens/s）。CUDA graph能够将多次内核启动合并为一次启动以加速推理过程。测试工作负载包含16个请求，每个请求输入长度1024，输出长度1024。
实验环境：QWen3-8B, 单卡TP1, H100 80GB GPU。禁用了基数缓存。
实验结果与分析：
- FlashInfer：禁用时吞吐量为441.73，启用后达到1245.51（2.82倍加速）。
- FA3：禁用时吞吐量为447.64，启用后达到1247.64（2.79倍加速）。
- Triton：禁用时吞吐量为419.64，启用后达到1228.36（2.93倍加速）。
- 结论：利用CUDA graph后，所有注意力内核均实现了至少2.79倍的推理速度提升。

实验3：离线推理性能测量

实验内容：使用三种常见的RL rollout工作负载（256个请求，输入/输出长度分别为1024/1024、4096/4096、8192/8192），测量非确定性模式和确定性模式的端到端延迟。
实验环境：QWen3-8B, 单卡TP1, H200 140GB GPU。禁用了基数缓存。
实验结果与分析：
- 确定性推理总体上是可用的，大部分减速范围在25%到45%之间。FlashInfer和FlashAttention 3后端的平均减速为34.35%。
- 以输入/输出1024为例，FlashInfer从30.85秒增加到43.99秒（+42.6%），FA3从34.70秒增加到44.14秒（+27.2%），Triton从36.91秒增加到57.25秒（+55.1%）。
- 结论：大部分开销来源于未优化的批处理不变内核（如矩阵乘法和注意力机制），这表明未来有巨大的性能优化空间。目前建议主要将该模式用于调试和复现研究。

补充细节

环境设置：为了设置环境，需要安装版本大于等于0.5.3的SGLang：

pip install "sglang[all]>=0.5.3"

启动服务器：SGLang支持跨多个模型的确定性推理。例如，对于Qwen3-8B，用户只需在启动服务器时添加--enable-deterministic-inference标志即可：

python3 -m sglang.launch_server \
    --model-path Qwen/Qwen3-8B \
    --attention-backend <flashinfer|fa3|triton> \
    --enable-deterministic-inference

结论

本文成功在SGLang中实现了稳定且可用的确定性推理，并通过与slime框架的深度集成，解锁了端到端100%可复现的强化学习训练流程。尽管目前的确定性推理在速度上相较于常规模式存在34.35%的性能损耗，但它为严谨的RL科学实验提供了必不可少的可靠性基石。

未来的工作将集中在以下关键领域以增强确定性推理：

加速批处理不变内核：优化并可能重写这些作为性能瓶颈的内核，目标是将性能差距缩小到20%以内，甚至与常规模式持平。
支持MoE模型：将确定性推理的支持范围从当前的密集模型（如QWen3-8B或LLaMa-3.1-8B）扩展到Qwen3-30B-A3B或DeepSeek-V3等混合专家（MoE）模型。
真正的同策略强化学习（True On-Policy RL）：进一步将确定性推理深度集成到强化学习框架（如slime）中实现可复现采样，最终达成真正的同策略训练。
增强基数缓存功能：改进基数树（radix tree）逻辑，使其不仅限于FlashAttention 3，还能兼容更广泛的注意力内核。
张量并行（Tensor Parallelism）支持：目前TP1和TP2因浮点加法顺序一致而具备确定性，未来将针对更大的TP设置修改归约内核以保证确定性。
FlexAttention集成：计划在未来扩展对FlexAttention后端的确定性推理支持。