[Sampling] Deterministic top-p / top-k renorm by default (fixes TP rank divergence in DFlash/DSpark)

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)

python/sglang/srt/layers/sampler.py

@@ -34,16 +34,20 @@
         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,
     )
 
 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,
     )
 

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"]

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

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
 

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__]))