fix(sglang): clamp deterministic sampler endpoints

relax/backends/sglang/deterministic_sampler_patch.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Backport SGLang's deterministic-sampler uint32 endpoint fix.
+
+SGLang 0.5.12.post1 maps a 32-bit hash to ``[0, 1]`` by dividing by
+``uint32.max``.  A hash equal to ``0xffffffff`` therefore produces exactly ``x
+== 1`` and Gumbel noise ``-log(-log(x)) == +inf``.  That token then wins the
+argmax regardless of its model probability.  Upstream clamps ``log(x)`` away
+from zero by one hash bucket; this module applies the same correction in the
+scheduler subprocess for the affected local runtime.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from importlib.metadata import version
+from inspect import signature
+from typing import Any
+
+import torch
+from packaging.version import Version
+
+from relax.utils.logging_utils import get_logger
+
+
+logger = get_logger(__name__)
+
+_AFFECTED_SGLANG_VERSIONS = frozenset({"0.5.12.post1"})
+_PATCH_MARKER = "_relax_uint32_endpoint_fix"
+
+
+def _installed_sglang_version() -> str:
+    return Version(version("sglang")).public
+
+
+def _uniform_hash_to_gumbel_(values: torch.Tensor) -> torch.Tensor:
+    """Transform uniform hash fractions in place without infinite endpoints."""
+    values.log_().clamp_(min=torch.finfo(values.dtype).min, max=-(2.0**-32)).neg_()
+    values.log_().neg_()
+    return values
+
+
+def _build_safe_multinomial_with_seed(
+    murmur_hash32: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
+    *,
+    compile_function: Callable[..., Any] = torch.compile,
+) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]:
+    """Build the upstream-equivalent deterministic multinomial function."""
+
+    def _safe_multinomial_with_seed(
+        logprobs: torch.Tensor, seed: torch.Tensor, positions: torch.Tensor
+    ) -> torch.Tensor:
+        _, vocabulary_size = logprobs.shape
+        seed = seed.to(torch.uint64)
+        column_indices = torch.arange(vocabulary_size, device=logprobs.device)
+        hashed = murmur_hash32(seed, positions, column_indices)
+        gumbel = hashed.to(torch.float64) / torch.iinfo(torch.uint32).max
+        _uniform_hash_to_gumbel_(gumbel)
+        gumbel.add_(logprobs.to(torch.float64))
+        return torch.argmax(gumbel, dim=1, keepdim=True)
+
+    return compile_function(dynamic=True)(_safe_multinomial_with_seed)
+
+
+def apply_deterministic_sampler_endpoint_patch() -> bool:
+    """Patch the affected SGLang sampler before scheduler model initialization.
+
+    Returns ``True`` only when this call installs the backport.  Unaffected
+    versions and already-patched scheduler processes are left unchanged.
+    """
+    installed_version = _installed_sglang_version()
+    if installed_version not in _AFFECTED_SGLANG_VERSIONS:
+        logger.info(
+            "SGLang deterministic sampler endpoint backport not required for version %s",
+            installed_version,
+        )
+        return False
+
+    from sglang.srt.layers import sampler
+    from sglang.srt.layers.utils.hash import murmur_hash32
+
+    current = sampler.multinomial_with_seed
+    if getattr(current, _PATCH_MARKER, False):
+        return False
+    parameter_names = tuple(signature(current).parameters)
+    if parameter_names != ("logprobs", "seed", "positions"):
+        raise RuntimeError(
+            "Affected SGLang multinomial_with_seed signature changed: "
+            f"expected=('logprobs', 'seed', 'positions'): actual={parameter_names}"
+        )
+
+    replacement = _build_safe_multinomial_with_seed(murmur_hash32)
+    setattr(replacement, _PATCH_MARKER, True)
+    sampler.multinomial_with_seed = replacement
+    logger.warning(
+        "Applied SGLang %s deterministic sampler uint32 endpoint backport",
+        installed_version,
+    )
+    return True

relax/backends/sglang/sglang_engine.py

@@ -162,16 +162,22 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int:
 
 
 def _patched_run_scheduler_process(*args, **kwargs):
-    """Scheduler-subprocess entry used for the routing-replay path.
+    """Scheduler-subprocess entry for Relax's SGLang runtime patches.
 
-    This wrapper is only installed when ``--optimize-routing-replay`` is
-    enabled (see ``_launch_server_with_patches``), so the routing-replay async
-    D→H patch is applied **unconditionally** here, preserving the original
-    behavior.
+    The deterministic sampler endpoint backport is version-gated internally.
+    The routing-replay async D→H patch remains gated by its existing runtime
+    environment flag.
     """
-    from relax.backends.sglang.routing_replay_patch import apply_patch
+    from relax.backends.sglang.deterministic_sampler_patch import (
+        apply_deterministic_sampler_endpoint_patch,
+    )
+
+    apply_deterministic_sampler_endpoint_patch()
+
+    if Envs.RELAX_OPTIMIZE_ROUTING_REPLAY:
+        from relax.backends.sglang.routing_replay_patch import apply_patch
 
-    apply_patch()
+        apply_patch()
 
     from sglang.srt.managers.scheduler import run_scheduler_process
 
@@ -184,9 +190,8 @@ def _launch_server_with_patches(server_args: ServerArgs):
 
     - main process: OPD pre-expanded multimodal patch
       (``RELAX_OPD_PREEXPANDED_PATCH=1``).
-    - scheduler subprocess: routing-replay (``RELAX_OPTIMIZE_ROUTING_REPLAY=1``)
-      installs ``_patched_run_scheduler_process``, which applies the
-      routing-replay patch unconditionally.
+    - scheduler subprocess: version-gated deterministic-sampler endpoint fix;
+      routing replay remains gated by ``RELAX_OPTIMIZE_ROUTING_REPLAY=1``.
     """
     from sglang.srt.entrypoints.http_server import launch_server
 
@@ -195,10 +200,7 @@ def _launch_server_with_patches(server_args: ServerArgs):
 
         apply_opd_preexpanded_patch()
 
-    if Envs.RELAX_OPTIMIZE_ROUTING_REPLAY:
-        launch_server(server_args, run_scheduler_process_func=_patched_run_scheduler_process)
-    else:
-        launch_server(server_args)
+    launch_server(server_args, run_scheduler_process_func=_patched_run_scheduler_process)
 
 
 def _resolve_external_model_arch(package_name):
@@ -234,9 +236,9 @@ def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
     multiprocessing.set_start_method("spawn", force=True)
     server_args.host = server_args.host.strip("[]")
 
-    # Each SGLang patch is controlled by its own env flag and applied
-    # independently (see ``_launch_server_with_patches`` and
-    # ``_patched_run_scheduler_process``); any combination is valid:
+    # Runtime patches are applied independently in the scheduler subprocess
+    # (see ``_launch_server_with_patches`` and ``_patched_run_scheduler_process``):
+    #   - deterministic sampler endpoint fix: version-gated backport
     #   - RELAX_OPTIMIZE_ROUTING_REPLAY : async D→H routing-replay patch (runtime)
     #   - RELAX_OPD_PREEXPANDED_PATCH   : OPD pre-expanded multimodal patch (runtime)
     #   - RELAX_OPD_PER_POS_TOKEN_IDS   : OPD per-position token_ids logprob;

tests/backends/sglang/test_deterministic_sampler_patch.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+
+import pytest
+import torch
+
+from relax.backends.sglang import deterministic_sampler_patch as patch
+
+
+def _identity_compile(*, dynamic):
+    assert dynamic is True
+    return lambda function: function
+
+
+def _install_fake_sglang_modules(
+    monkeypatch: pytest.MonkeyPatch,
+    sampler: ModuleType,
+    hash_module: ModuleType,
+) -> None:
+    sglang = ModuleType("sglang")
+    srt = ModuleType("sglang.srt")
+    layers = ModuleType("sglang.srt.layers")
+    utils = ModuleType("sglang.srt.layers.utils")
+    sglang.srt = srt
+    srt.layers = layers
+    layers.sampler = sampler
+    layers.utils = utils
+    utils.hash = hash_module
+    monkeypatch.setitem(sys.modules, "sglang", sglang)
+    monkeypatch.setitem(sys.modules, "sglang.srt", srt)
+    monkeypatch.setitem(sys.modules, "sglang.srt.layers", layers)
+    monkeypatch.setitem(sys.modules, "sglang.srt.layers.sampler", sampler)
+    monkeypatch.setitem(sys.modules, "sglang.srt.layers.utils", utils)
+    monkeypatch.setitem(sys.modules, "sglang.srt.layers.utils.hash", hash_module)
+
+
+def test_uniform_hash_endpoint_has_finite_upstream_gumbel_cap():
+    values = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float64)
+
+    result = patch._uniform_hash_to_gumbel_(values)
+
+    assert torch.isfinite(result).all()
+    assert result[1].item() == pytest.approx(-torch.log(-torch.log(torch.tensor(0.5))).item())
+    assert result[-1].item() == pytest.approx(-torch.log(torch.tensor(2.0**-32)).item())
+
+
+def test_safe_multinomial_does_not_let_uint32_endpoint_override_logprob():
+    uint32_max = torch.iinfo(torch.uint32).max
+
+    def fake_hash(seed, positions, column_indices):
+        assert seed.shape == positions.shape == (1,)
+        assert column_indices.shape == (2,)
+        return torch.tensor([[uint32_max, uint32_max // 2]], dtype=torch.uint32)
+
+    sample = patch._build_safe_multinomial_with_seed(fake_hash, compile_function=_identity_compile)
+    selected = sample(
+        torch.tensor([[float("-inf"), 0.0]], dtype=torch.float64),
+        torch.tensor([44], dtype=torch.int64),
+        torch.tensor([179], dtype=torch.int64),
+    )
+
+    assert selected.tolist() == [[1]]
+
+
+def test_apply_patch_is_version_gated_and_idempotent(monkeypatch):
+    sampler = ModuleType("sglang.srt.layers.sampler")
+
+    def original(logprobs, seed, positions):
+        return logprobs, seed, positions
+
+    sampler.multinomial_with_seed = original
+    hash_module = ModuleType("sglang.srt.layers.utils.hash")
+    hash_module.murmur_hash32 = object()
+    _install_fake_sglang_modules(monkeypatch, sampler, hash_module)
+    monkeypatch.setattr(patch, "_installed_sglang_version", lambda: "0.5.12.post1")
+    replacement = lambda *_args: None
+    monkeypatch.setattr(patch, "_build_safe_multinomial_with_seed", lambda _hash: replacement)
+
+    assert patch.apply_deterministic_sampler_endpoint_patch() is True
+    assert sampler.multinomial_with_seed is replacement
+    assert getattr(replacement, patch._PATCH_MARKER) is True
+    assert patch.apply_deterministic_sampler_endpoint_patch() is False
+
+
+def test_apply_patch_leaves_unaffected_sglang_unchanged(monkeypatch):
+    monkeypatch.setattr(patch, "_installed_sglang_version", lambda: "0.5.13")
+
+    assert patch.apply_deterministic_sampler_endpoint_patch() is False
+
+
+def test_local_version_suffix_resolves_to_affected_public_version(monkeypatch):
+    monkeypatch.setattr(patch, "version", lambda _package: "0.5.12.post1+cu129")
+
+    assert patch._installed_sglang_version() == "0.5.12.post1"
+
+
+def test_affected_signature_drift_fails_closed(monkeypatch):
+    sampler = ModuleType("sglang.srt.layers.sampler")
+    sampler.multinomial_with_seed = lambda inputs, seed: (inputs, seed)
+    hash_module = ModuleType("sglang.srt.layers.utils.hash")
+    hash_module.murmur_hash32 = object()
+    _install_fake_sglang_modules(monkeypatch, sampler, hash_module)
+    monkeypatch.setattr(patch, "_installed_sglang_version", lambda: "0.5.12.post1")
+
+    with pytest.raises(RuntimeError, match="signature changed"):
+        patch.apply_deterministic_sampler_endpoint_patch()

tests/backends/sglang/test_router_registration.py

@@ -2,6 +2,7 @@
 
 import importlib
 import logging
+import pickle
 import sys
 from types import ModuleType, SimpleNamespace
 
@@ -148,6 +149,72 @@ def _make_engine(sglang_engine_module):
     return engine
 
 
+@pytest.mark.parametrize("routing_replay", [False, True])
+def test_scheduler_wrapper_applies_endpoint_patch_before_optional_routing_patch(
+    monkeypatch, sglang_engine_module, routing_replay
+):
+    events = []
+    endpoint_patch_module = ModuleType("relax.backends.sglang.deterministic_sampler_patch")
+    endpoint_patch_module.apply_deterministic_sampler_endpoint_patch = lambda: events.append("endpoint")
+    routing_patch_module = ModuleType("relax.backends.sglang.routing_replay_patch")
+    routing_patch_module.apply_patch = lambda: events.append("routing")
+    scheduler_module = ModuleType("sglang.srt.managers.scheduler")
+
+    def run_scheduler_process(*args, **kwargs):
+        events.append(("scheduler", args, kwargs))
+        return "finished"
+
+    scheduler_module.run_scheduler_process = run_scheduler_process
+    monkeypatch.setitem(sys.modules, endpoint_patch_module.__name__, endpoint_patch_module)
+    monkeypatch.setitem(sys.modules, routing_patch_module.__name__, routing_patch_module)
+    monkeypatch.setitem(sys.modules, scheduler_module.__name__, scheduler_module)
+    monkeypatch.setattr(
+        sglang_engine_module.Envs,
+        "RELAX_OPTIMIZE_ROUTING_REPLAY",
+        routing_replay,
+        raising=False,
+    )
+
+    assert sglang_engine_module._patched_run_scheduler_process("arg", key="value") == "finished"
+    expected = ["endpoint"]
+    if routing_replay:
+        expected.append("routing")
+    assert events[:-1] == expected
+    assert events[-1] == (("scheduler", ("arg",), {"key": "value"}))
+
+
+@pytest.mark.parametrize("routing_replay", [False, True])
+def test_launch_server_always_receives_picklable_scheduler_wrapper(monkeypatch, sglang_engine_module, routing_replay):
+    calls = []
+    http_server = ModuleType("sglang.srt.entrypoints.http_server")
+
+    def launch_server(server_args, **kwargs):
+        calls.append((server_args, kwargs))
+
+    http_server.launch_server = launch_server
+    monkeypatch.setitem(sys.modules, http_server.__name__, http_server)
+    monkeypatch.setattr(sglang_engine_module.Envs, "RELAX_OPD_PREEXPANDED_PATCH", False, raising=False)
+    monkeypatch.setattr(
+        sglang_engine_module.Envs,
+        "RELAX_OPTIMIZE_ROUTING_REPLAY",
+        routing_replay,
+        raising=False,
+    )
+    server_args = object()
+
+    sglang_engine_module._launch_server_with_patches(server_args)
+
+    assert calls == [
+        (
+            server_args,
+            {"run_scheduler_process_func": sglang_engine_module._patched_run_scheduler_process},
+        )
+    ]
+    assert pickle.loads(pickle.dumps(sglang_engine_module._patched_run_scheduler_process)).__name__ == (
+        "_patched_run_scheduler_process"
+    )
+
+
 def test_missing_load_format_choices_uses_legacy_remote(sglang_engine_module):
     assert sglang_engine_module._preferred_s3_stream_load_format() == "remote"