SGLang 实现确定性推理，并与 slime 合作实现 100% 可复现 RL 训练

TL;DR：本文分享我们在 SGLang 中实现确定性推理的努力，以及与 slime 合作迈向可复现 RL 训练的进展。

最近，Thinking Machines Lab 发布了一篇博客，详细介绍了他们的发现。该博客发布后，业界反响热烈，急切期待开源推理引擎能够实现稳定可用的确定性推理，甚至更进一步，实现完全可复现的 RL 训练。现在，SGLang 和 slime 共同给出了答案。

在 Thinking Machines Lab 的批次不变算子基础上，SGLang 实现了完全确定性推理，同时保持与 chunked prefill、CUDA graphs、radix cache 和 非贪心采样的兼容性。借助 CUDA graphs，SGLang 实现了 2.8 倍加速，并将性能开销降至仅 34.35%（相比 TML 的 61.5%）。

在此基础上，SGLang 与 slime 团队合作，进一步解锁了 100% 可复现的 RL 训练——这一突破仅需极少的代码改动。我们在 Qwen3-8B 上的验证实验展示了完美的可复现性：两次独立训练运行产生完全相同的曲线，为严谨的科学实验提供了所需的可靠性。

可复现指南

现在让我们深入了解一些技术细节。

为什么确定性推理很重要

从大语言模型（LLM）推理中获得一致输出的能力正变得越来越重要。例如，正如研究者指出，推理结果的不确定性会隐式地将同策略强化学习（RL）转变为异策略 RL。然而，即使我们在 SGLang 中将温度降至 0，由于使用了动态批处理和 radix cache，采样仍然不是确定性的（过往讨论见此处）。

正如博客中所述，不确定性的最大来源是变化的批大小：即使用户反复提交相同的提示，输出也可能因运行而异，因为该请求可能与其他用户的请求一起批处理，而批大小的差异会导致不确定的推理结果。

进一步解释，不同的批大小会影响内核的归约拆分过程。这导致每个归约块的顺序和大小发生变化，由于浮点运算的非结合性，可能造成不确定的输出。为解决此问题，他们将归约内核（RMSNorm、矩阵乘法、注意力等）替换为批次不变的实现。这些内核也作为配套库发布，供外部集成使用。

He, Horace and Thinking Machines Lab, "Defeating Nondeterminism in LLM Inference", Thinking Machines Lab: Connectionism, Sep 2025.

在 Thinking Machines Lab 工作的基础上，SGLang 为确定性 LLM 推理提供了稳健、高吞吐的解决方案，将批次不变内核、CUDA graphs、radix cache 和 chunked prefill 与高效性能相结合。确定性已通过全面的测试和 RL 训练实验得到广泛验证。

关键增强包括：

• 集成来自 Thinking Machines Lab 的批次不变内核，包括均值、log-softmax 和矩阵乘法内核。
• 实现批次不变注意力内核，采用固定的 split-KV 大小。支持多种后端，包括 FlashInfer、FlashAttention 3 和 Triton。
• 完全兼容常见推理特性，如分块预填充、CUDA graph、基数缓存，在启用确定性推理时这些特性均保持支持。
• 在采样参数中暴露每请求种子，允许用户在 temperature > 0 时也能启用确定性推理。
• 更优的性能：与 TML 博客中报告的 61.5% 减速相比，SGLang 在使用 FlashInfer 和 FlashAttention 3 后端时平均减速仅为 34.35%，实现了显著提升。使用 CUDA graph 时，相比最小集成可实现 2.8 倍加速。

结果

验证确定性行为

我们引入一个确定性测试套件，用于验证推理结果在不同批次条件下是否保持一致。该测试涵盖三个子测试，从简单到更具挑战性递进：

• 单一：在不同批次大小下运行相同提示，检查输出是否保持相同。
• 混合：在同一批次中混合不同类型的提示（短提示和长提示），验证一致性。
• 前缀：使用源自同一长文本但前缀长度不同的提示，随机分批，测试结果在多次试验中是否可复现。

以下是 50 次采样试验的结果。数字表示每个子测试中观察到的唯一输出数量（越低越确定）。

*在 QWen3-8B 上测试

* 启用了 Cuda graph、分块预填充。FlashInfer 和 Triton 的基数缓存被禁用，因为其支持仍在进行中。

CUDA Graph 加速

CUDA graph 可以通过将多个内核启动合并为单次启动来加速推理过程。我们的评估比较了启用和不启用 CUDA graph 时确定性推理的总吞吐量。测试工作负载包含 16 个请求，每个请求输入长度为 1024，输出长度为 1024。结果显示，使用 CUDA graph 时，所有注意力内核至少实现 2.79 倍加速。

*设置：QWen3-8B，TP1，H100 80GB

*我们在所有性能基准测试中禁用了基数缓存，因为 FlashInfer 和 Triton 的基数缓存支持仍在进行中。

测量离线推理性能

我们使用三种常见的 RL rollout 工作负载（256 个请求，输入/输出长度各异）测量了非确定性和确定性模式的端到端延迟。

确定性推理总体可用，大多数减速在 25% 到 45% 之间，FlashInfer 和 FlashAttention 3 后端的平均减速为 34.35%。大部分开销来自未优化的批次不变内核（矩阵乘法和注意力），表明性能仍有显著提升空间。

*设置：QWen3-8B，TP1，H200 140GB。

*我们在所有性能基准测试中禁用了基数缓存，因为 FlashInfer 和 Triton 的基数缓存支持仍在开发中。

我们承认确定性推理比普通模式慢得多。我们建议主要将其用于调试和可复现性。未来的工作将专注于加速确定性推理，目标是将性能差距缩小到 20% 以下，或理想情况下与普通模式持平。

使用方法

环境设置

要设置环境，请安装 SGLang 版本 >=0.5.3

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 的分块预填充技术旨在管理长上下文请求。然而，其默认的分块策略违反了注意力内核的确定性要求。

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

标准分块策略基于“尽力而为”原则运行。在此示例中，该策略尝试通过将 seq_b 的 b2 单元拆分为两个较小的部分来生成 8,192 个 token 的 chunk_1。这可能导致不一致的截断点，因为拆分后 b2 的长度取决于 seq_a 的长度。为了解决这个问题，我们调整了分块逻辑，使截断点与 split_kv_size 的整数倍对齐。此调整确保 b2 的处理推迟到后续块，使其能够作为完整单元由注意力内核计算。

注意力后端

注意力内核是确定性的重要组成部分。对于不同的注意力后端，我们以不同方式修改它们以满足其使用要求。

• 对于 Flashinfer 后端，我们利用 batch invariant FA2 kernels 中的 fixed_split_size 和 disable_kv_split 参数来在内核规划期间固定拆分大小。分块预填充的截断与预填充拆分大小对齐。（PR 链接）
• 对于 FlashAttention-3 后端，flash attention 内核的 num-splits 固定为 1 以确保确定性。（PR 链接）
• 对于 Triton 后端，我们固定解码的拆分大小，并手动设置分块预填充的对齐大小。借助 Triton 后端的可扩展性，确定性推理也可以在 AMD 硬件上运行。（PR 链接）。

可复现的非贪婪采样

为了将确定性扩展到贪婪解码之外，我们引入了一个新的采样函数：multinomial_with_seed。

该算子不依赖于在批处理下本质上不确定的 torch.multinomial，而是使用由种子哈希函数生成的 Gumbel 噪声来扰动 logits。因此，相同的 (inputs, seed) 对总是产生相同的样本，即使 temperature > 0。

这一修改实现了确定性多项采样，同时保留了强化学习 rollout 所需的随机性。

RL 框架集成（slime）

我们将温度 > 0 的确定性推理集成到了 slime 的 GRPO 训练配方中。在初步实验中，重复的 RL 训练运行在前几次迭代中产生了完全相同的 rollout 响应和损失值，证实了 rollout 过程本身是确定性的。

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

• Flash Attention：使用 Flash Attention v2 而非 v3，以实现确定性的反向传播
• Megatron：设置 --deterministic-mode 标志以进行确定性训练
• Environment Variables: Configure critical settings:
  • NCCL_ALGO=Ring
  • NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
  • CUBLAS_WORKSPACE_CONFIG=:4096:8
• PyTorch：启用 torch.use_deterministic_algorithms(True, warn_only=False)

通过这些全面的改动，我们成功实现了 slime 中 GRPO 的完整训练可复现性，从而实现了真正确定性的端到端 RL 训练流水线。

未来工作

我们未来的工作将聚焦于通过解决以下关键领域来增强确定性推理：

• 更快的批次不变内核：批次不变内核是性能瓶颈，因此我们将致力于优化其配置，并可能重写它们以提升性能。这对于提高 RL rollout 的速度也至关重要。
• 支持 MoE 模型：目前我们仅支持 QWen3-8B 或 LLaMa-3.1-8B 等稠密模型的确定性推理。未来我们计划将支持扩展到 Qwen3-30B-A3B 或 DeepSeek-V3 等 MoE 模型。
• 真正的 On-Policy RL：我们计划进一步将确定性推理集成到强化学习框架（例如 slime）中，以实现可复现的采样，最终目标是实现真正的 on-policy 训练。
• 增强 Radix Cache 功能：我们将改进基数树，使其能够兼容更多种类的注意力内核，突破当前仅限 FlashAttention 3 后端的限制。
• 张量并行：TP1 和 TP2 由于浮点加法顺序一致而是确定性的；更大的 TP 配置需要修改归约内核以实现确定性。
• FlexAttention 集成：除了当前支持的注意力后端外，我们计划在未来将确定性推理的支持扩展到 FlexAttention。
• 确定性推理功能的路线图可在此 issue 中找到。

SGLang 的确定性推理和 slime 的可复现训练能力目前正在积极开发和改进中。我们诚挚欢迎用户和开发者积极尝试这些功能，并向我们提供宝贵的反馈。您的经验和建议将帮助我们进一步优化这些重要能力，并推动确定性推理技术的发展。

致谢

我们衷心感谢以下团队和合作者：

• SGLang 团队和社区：Baizhou Zhang、Biao He、Qiaolin Yu、Xinyuan Tong、Ke Bao、Yineng Zhang、Chi Zhang、Ying Sheng、Lianmin Zheng 以及许多其他人
• Flashinfer 团队和社区：Wenxuan Tan、Yilong Zhao、Zihao Ye
• slime 团队和社区：Zilin Zhu
• AMD：Yusheng Su
• Thinking Machines Lab：感谢他们出色的 博客 和 batch_invariant_ops 库