From e8ee82d42fd050dd377176331157f19f9322bfdc Mon Sep 17 00:00:00 2001 From: Gilford Ting Date: Tue, 8 Sep 2026 21:08:34 +0000 Subject: [PATCH] [Sampling] Deterministic top-p / top-k renorm by default (fixes TP rank divergence in DFlash/DSpark) flashinfer's default top_p_renorm_probs (AIR radix, >= 0.6.7) and top_k_renorm_probs (radix multi-CTA) accumulate float sums with atomicAdd, so two calls on byte-identical input return probabilities that differ in the last bits. Every TP rank runs them independently on the same logits in speculative verification (DFlash, DSpark, EAGLE) and in the sampler's min_p path; a last-bit gap flips a rejection-sampling coin or the bonus token on one rank only, the per-rank radix/KV caches drift, and a later prefix match deadlocks an NCCL collective (#33549, #33289; #33614 is the broadcast workaround). Add sglang.srt.layers.sampling_renorm with top_p_renorm_prob / top_k_renorm_prob that default to deterministic kernels: flashinfer's integer-histogram AIR (is_deterministic=True) for top-p, and the single-CTA kernel already compiled into sgl_kernel for top-k. Route the three call sites (sampler, dflash_utils, eagle_utils) through it. New `deterministic` kwarg and SGLANG_RENORM_DETERMINISTIC env var opt back into the faster kernels. Add regression tests asserting bit-identical output across repeated calls. The change lives in sglang rather than the sgl_kernel wrappers because sglang-kernel ships as a pinned prebuilt wheel; a wrapper change would not be testable in CI or reach users until the next kernel release. Measured: TP=2 DFlash with top_p=0.9 and 32 streams wedged within 4 min on H100 and B300 three times out of three; with deterministic renorm it ran 30 min, 21k requests, zero cross-rank divergence. Co-Authored-By: Claude Fable 5.1 --- python/sglang/srt/environ.py | 8 ++ python/sglang/srt/layers/sampler.py | 8 +- python/sglang/srt/layers/sampling_renorm.py | 127 ++++++++++++++++++ python/sglang/srt/speculative/dflash_utils.py | 5 +- python/sglang/srt/speculative/eagle_utils.py | 10 +- .../registered/kernels/ops/moe/test_renorm.py | 73 +++++++++- 6 files changed, 220 insertions(+), 11 deletions(-) create mode 100644 python/sglang/srt/layers/sampling_renorm.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6f64ca7956c0..bde801719b5a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1230,6 +1230,14 @@ class Envs: # Deterministic inference and all-reduce # =================================================================== SGLANG_ENABLE_DETERMINISTIC_INFERENCE = EnvBool(False) + # Used by sglang.srt.layers.sampling_renorm: when true (default) the top-p / + # top-k renormalization used by the sampler and speculative verification runs + # on kernels whose output is bit-identical call to call. The flashinfer defaults accumulate with float + # atomics and can differ in the last bits between TP ranks running the same + # input, which desynchronizes committed tokens / speculative accept lengths + # across ranks and eventually deadlocks a collective (#33549, #33289). Set to + # 0 to opt back into the faster non-deterministic kernels. + SGLANG_RENORM_DETERMINISTIC = EnvBool(True) # Use 1-stage all-reduce kernel on AMD (deterministic, fixed accumulation order) # If not set: auto (enabled when --enable-deterministic-inference is on) # Set to 1: force enable (even without --enable-deterministic-inference) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 7781f608b306..a8562b3e03ac 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -34,7 +34,8 @@ min_p_sampling_from_probs, top_k_top_p_sampling_from_probs, ) - from sgl_kernel import ( + + from sglang.srt.layers.sampling_renorm import ( top_k_renorm_prob, top_p_renorm_prob, ) @@ -42,8 +43,11 @@ if is_musa(): from sgl_kernel import ( min_p_sampling_from_probs, - top_k_renorm_prob, top_k_top_p_sampling_from_probs, + ) + + from sglang.srt.layers.sampling_renorm import ( + top_k_renorm_prob, top_p_renorm_prob, ) diff --git a/python/sglang/srt/layers/sampling_renorm.py b/python/sglang/srt/layers/sampling_renorm.py new file mode 100644 index 000000000000..565e68331780 --- /dev/null +++ b/python/sglang/srt/layers/sampling_renorm.py @@ -0,0 +1,127 @@ +"""Top-p / top-k probability renormalization with deterministic output. + +Why this module exists +---------------------- +flashinfer's default ``top_p_renorm_probs`` (AIR radix, flashinfer >= 0.6.7) +and ``top_k_renorm_probs`` (radix multi-CTA) pool partial sums across thread +blocks with float ``atomicAdd``. Float addition is not associative, so two +calls on byte-identical input can return probabilities that differ in the +last bits (measured: 99 of 99 repeated calls differ on a 256 x 128256 flat +distribution; for top-k even on peaky rows). + +Every tensor-parallel rank runs these kernels independently on the same +logits, and the output feeds decisions that are committed to per-rank state: +sampled tokens (the ``min_p`` path of the sampler) and speculative-decoding +accept lengths / bonus tokens (DFlash, DSpark, EAGLE verify). A last-bit gap +between ranks flips a rejection-sampling coin on one rank only, the per-rank +radix/KV caches drift apart, and a later prefix match deadlocks an NCCL +collective (#33549, #33289; #33614 is the broadcast that papers over it). + +So by default this module routes to kernels whose output is bit-identical +call to call: + +* top-p: flashinfer's integer-histogram AIR variant (``is_deterministic=True``); + without flashinfer (MUSA), the single-CTA kernel compiled into ``sgl_kernel``. +* top-k: the single-CTA kernel compiled into ``sgl_kernel`` (fixed-order block + reductions). flashinfer has no deterministic option for its radix top-k. + +Set ``SGLANG_RENORM_DETERMINISTIC=0`` (or pass ``deterministic=False``) to opt +back into the faster non-deterministic kernels, e.g. on a single rank without +speculative decoding. +""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch + +from sglang.srt.environ import envs + +try: + import flashinfer.sampling as _flashinfer_sampling + + _HAS_FLASHINFER = True +except ImportError: # pragma: no cover - non-CUDA builds + _flashinfer_sampling = None + _HAS_FLASHINFER = False + +import sgl_kernel as _sgl_kernel + + +def _resolve(deterministic: Optional[bool]) -> bool: + if deterministic is None: + return envs.SGLANG_RENORM_DETERMINISTIC.get() + return bool(deterministic) + + +def _split_param(x: Union[torch.Tensor, float, int]): + """(per-row tensor or None, scalar) for the sgl_kernel op schemas.""" + if isinstance(x, torch.Tensor): + return x, 0 + return None, x + + +def _single_cta_top_k( + probs: torch.Tensor, top_k: Union[torch.Tensor, int] +) -> torch.Tensor: + # torch.ops.sgl_kernel.top_k_renorm_probs is flashinfer's single-CTA kernel + # compiled into sgl_kernel: one block per row, fixed-order block reductions. + arr, val = _split_param(top_k) + probs32 = probs.float() + out = torch.empty_like(probs32) + torch.ops.sgl_kernel.top_k_renorm_probs.default( + probs32, out, arr.int() if arr is not None else None, int(val) + ) + return out if out.dtype == probs.dtype else out.to(probs.dtype) + + +def _single_cta_top_p( + probs: torch.Tensor, top_p: Union[torch.Tensor, float] +) -> torch.Tensor: + arr, val = _split_param(top_p) + probs32 = probs.float() + out = torch.empty_like(probs32) + torch.ops.sgl_kernel.top_p_renorm_probs.default( + probs32, out, arr.float() if arr is not None else None, float(val) + ) + return out if out.dtype == probs.dtype else out.to(probs.dtype) + + +def top_p_renorm_prob( + probs: torch.Tensor, + top_p: Union[torch.Tensor, float], + deterministic: Optional[bool] = None, +) -> torch.Tensor: + """Zero every token outside the top-p nucleus and renormalize. + + ``deterministic`` defaults to ``SGLANG_RENORM_DETERMINISTIC`` (on). When on, + repeated calls on identical input return bit-identical output. + """ + if not _resolve(deterministic): + return _sgl_kernel.top_p_renorm_prob(probs, top_p) + if _HAS_FLASHINFER and probs.is_cuda: + # is_deterministic exists since flashinfer 0.6.7; sglang pins newer. + return _flashinfer_sampling.top_p_renorm_probs( + probs, top_p, is_deterministic=True + ) + # no flashinfer (e.g. MUSA): the single-CTA kernel is deterministic already. + return _single_cta_top_p(probs, top_p) + + +def top_k_renorm_prob( + probs: torch.Tensor, + top_k: Union[torch.Tensor, int], + deterministic: Optional[bool] = None, +) -> torch.Tensor: + """Zero every token outside the top-k set and renormalize. + + ``deterministic`` defaults to ``SGLANG_RENORM_DETERMINISTIC`` (on). When on, + repeated calls on identical input return bit-identical output. + """ + if not _resolve(deterministic): + return _sgl_kernel.top_k_renorm_prob(probs, top_k) + return _single_cta_top_k(probs, top_k) + + +__all__ = ["top_k_renorm_prob", "top_p_renorm_prob"] diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 522790a5380e..2733376223c0 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -42,10 +42,11 @@ if is_cuda() or is_musa(): try: - from sgl_kernel import ( + from sgl_kernel import tree_speculative_sampling_target_only + + from sglang.srt.layers.sampling_renorm import ( top_k_renorm_prob, top_p_renorm_prob, - tree_speculative_sampling_target_only, ) _DFLASH_SAMPLING_VERIFY_AVAILABLE = True diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 616f808d7ba6..3c709659f87f 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -855,15 +855,15 @@ def eagle_sample( tree_speculative_sampling_target_only, ) else: - from sgl_kernel import ( - top_k_renorm_prob, - top_p_renorm_prob, - tree_speculative_sampling_target_only, - ) + from sgl_kernel import tree_speculative_sampling_target_only from sglang.kernels.ops.speculative.reject_sampling import ( chain_speculative_sampling_triton, ) + from sglang.srt.layers.sampling_renorm import ( + top_k_renorm_prob, + top_p_renorm_prob, + ) use_rejection_sampling = get_spec().speculative_use_rejection_sampling diff --git a/test/registered/kernels/ops/moe/test_renorm.py b/test/registered/kernels/ops/moe/test_renorm.py index ae66e0ff9881..1395be79e078 100644 --- a/test/registered/kernels/ops/moe/test_renorm.py +++ b/test/registered/kernels/ops/moe/test_renorm.py @@ -1,6 +1,7 @@ # Adapted from https://github.com/flashinfer-ai/flashinfer/blob/main/tests/test_sampling.py # and /sgl-workspace/sglang/python/sglang/kernels/aot/tests/test_sampling.py +import os import sys import pytest @@ -9,7 +10,7 @@ from sglang.srt.utils import is_hip from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x") if is_hip(): @@ -20,7 +21,10 @@ top_p_renorm_probs_triton as top_p_renorm_prob, ) else: - from sgl_kernel import top_k_renorm_prob, top_p_renorm_prob + from sglang.srt.layers.sampling_renorm import ( + top_k_renorm_prob, + top_p_renorm_prob, + ) @pytest.mark.parametrize("batch_size", [1, 99, 989]) @@ -92,5 +96,70 @@ def test_top_p_renorm_probs(batch_size, vocab_size, p): ) +def _default_is_deterministic() -> bool: + return os.environ.get("SGLANG_RENORM_DETERMINISTIC", "1").lower() not in ( + "0", + "false", + "no", + "off", + ) + + +def _flat_distribution(batch_size: int, vocab_size: int) -> torch.Tensor: + """A heavy-tailed, high-entropy distribution: thousands of tokens survive a + top-p 0.95 cutoff, so the renorm kernels' histogram sums have many terms. + This is the regime in which the flashinfer float-atomic kernels return + different bytes call to call (see SGLANG_RENORM_DETERMINISTIC).""" + gen = torch.Generator(device="cuda:0").manual_seed(7) + ranks = torch.arange(1, vocab_size + 1, device="cuda:0", dtype=torch.float32) + zipf = -1.1 * torch.log(ranks) + logits = torch.empty(batch_size, vocab_size, device="cuda:0") + for r in range(batch_size): + perm = torch.randperm(vocab_size, device="cuda:0", generator=gen) + logits[r] = zipf[perm] + 0.6 * torch.randn( + vocab_size, device="cuda:0", generator=gen + ) + return torch.softmax(logits, dim=-1) + + +@pytest.mark.skipif(is_hip(), reason="CUDA-only: exercises the flashinfer dispatch") +@pytest.mark.parametrize("batch_size,vocab_size", [(256, 128256)]) +def test_top_p_renorm_probs_is_deterministic(batch_size, vocab_size): + """Repeated calls on the same input must return bit-identical output. + + Every TP rank runs this kernel independently on the same logits, and its + output feeds sampled tokens and speculative accept decisions that are + committed to per-rank KV/radix state. A last-bit difference between ranks + desynchronizes them and eventually deadlocks a collective (#33549, #33289). + """ + probs = _flat_distribution(batch_size, vocab_size) + top_p = torch.full((batch_size,), 0.95, device="cuda:0") + ref = top_p_renorm_prob(probs, top_p, deterministic=True) + for _ in range(30): + out = top_p_renorm_prob(probs, top_p, deterministic=True) + assert torch.equal(out, ref), "top_p_renorm_prob output changed between calls" + # the default path is the deterministic one unless the env opts out + if _default_is_deterministic(): + assert torch.equal(top_p_renorm_prob(probs, top_p), ref) + + +@pytest.mark.skipif(is_hip(), reason="CUDA-only: exercises the flashinfer dispatch") +@pytest.mark.parametrize("batch_size,vocab_size,k", [(256, 128256, 40)]) +def test_top_k_renorm_probs_is_deterministic(batch_size, vocab_size, k): + """Same invariant for top-k renorm; flashinfer's radix multi-CTA kernel + accumulates the kept mass with float atomics and is not reproducible.""" + probs = _flat_distribution(batch_size, vocab_size) + top_k = torch.full((batch_size,), k, device="cuda:0", dtype=torch.int32) + ref = top_k_renorm_prob(probs, top_k, deterministic=True) + for _ in range(30): + out = top_k_renorm_prob(probs, top_k, deterministic=True) + assert torch.equal(out, ref), "top_k_renorm_prob output changed between calls" + if _default_is_deterministic(): + assert torch.equal(top_k_renorm_prob(probs, top_k), ref) + # and the deterministic kernel must agree with the fast one up to rounding + fast = top_k_renorm_prob(probs, top_k, deterministic=False) + torch.testing.assert_close(fast, ref, rtol=1e-5, atol=1e-6) + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))