diff --git a/relax/backends/sglang/deterministic_sampler_patch.py b/relax/backends/sglang/deterministic_sampler_patch.py new file mode 100644 index 000000000..1226053ca --- /dev/null +++ b/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 diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 15ca1d172..b167f73a3 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/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; diff --git a/tests/backends/sglang/test_deterministic_sampler_patch.py b/tests/backends/sglang/test_deterministic_sampler_patch.py new file mode 100644 index 000000000..75ab6add9 --- /dev/null +++ b/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() diff --git a/tests/backends/sglang/test_router_registration.py b/tests/backends/sglang/test_router_registration.py index 84cc8f3e0..201788bc1 100644 --- a/tests/backends/sglang/test_router_registration.py +++ b/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"