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