From 5c95bf15c83b5725792bec79df7283b8905f5c6f Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 4 Jul 2026 02:22:05 -0700 Subject: [PATCH] [refactor] Add the post-process resolution stage; migrate sampling_backend (stack 10/15) (#30072) Co-authored-by: Claude Fable 5 --- python/sglang/srt/arg_groups/overrides.py | 96 ++++++++++++++++- python/sglang/srt/runtime_context.py | 1 + python/sglang/srt/server_args.py | 28 +++-- test/registered/unit/test_model_overrides.py | 103 +++++++++++++++++-- 4 files changed, 210 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index b08cef78a..766064f93 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -34,7 +34,7 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tupl from sglang.srt.arg_groups.arg_utils import model_overridable_fields from sglang.srt.runtime_context import resolve_flag_leaf -from sglang.srt.utils.common import is_xpu +from sglang.srt.utils.common import is_flashinfer_available, is_xpu logger = logging.getLogger(__name__) @@ -95,6 +95,72 @@ def _invoke_provider( return declared +class ResolvedView: + """Read-only view of the resolving configuration handed to post-process + passes. + + During the dual-apply transition the view forwards every read to the live + ``server_args`` — the pristine input plus the declarations replayed so far + plus any residual imperative writes — which is exactly the state the + legacy handler at the same slot observed. In the end state (dual-apply + retired) the same type overlays the accumulated declarations on the + pristine object. Writes are rejected: passes return declarations. + """ + + __slots__ = ("_server_args", "_overlay") + + def __init__(self, server_args: Any, overlay: Optional[Dict[str, Any]] = None): + object.__setattr__(self, "_server_args", server_args) + object.__setattr__(self, "_overlay", overlay or {}) + + def __getattr__(self, name: str) -> Any: + overlay = object.__getattribute__(self, "_overlay") + if name in overlay: + return overlay[name] + return getattr(object.__getattribute__(self, "_server_args"), name) + + def __setattr__(self, name: str, value: Any) -> None: + raise AttributeError( + "ResolvedView is read-only; post-process passes return declarations" + ) + + +# Ordered post-process passes (the normalization stage). List order is the +# end-state execution order and mirrors today's handler call sequence in +# __post_init__; during the transition each pass is invoked from its legacy +# slot via run_post_process_pass, so ordering is preserved byte-for-byte. +POST_PROCESS_PASSES: List[Callable[..., dict]] = [] + + +def register_post_process(fn: Callable[..., dict]) -> Callable[..., dict]: + """Register a post-process pass: ``fn(view) -> {field: resolved_value}``. + + The pass reads a :class:`ResolvedView` (post-model-override state) and + must not mutate anything; validations may live in a pass (read + raise). + """ + POST_PROCESS_PASSES.append(fn) + return fn + + +def run_post_process_pass(server_args: Any, fn: Callable[..., dict]) -> None: + """Transition-period invocation of one pass at its legacy handler slot. + + Evaluates the pass on the live state (through a read-only view), appends + its declaration to the declaration stash, and dual-applies it in place — + byte-identical to the imperative handler write this replaces. + """ + declared = fn(ResolvedView(server_args)) + if not isinstance(declared, dict): + raise TypeError( + f"post-process pass {fn.__qualname__} must return a dict, " + f"got {type(declared).__name__}" + ) + if declared: + entry = (fn.__qualname__, dict(declared)) + server_args._resolved_overrides.append(entry) + apply_declarations_to_server_args(server_args, [entry]) + + def collect_model_override_declarations( architecture: str, server_args: Any, hf_config: Any ) -> List[Tuple[str, Dict[str, Any]]]: @@ -239,6 +305,34 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: return overrides +# --------------------------------------------------------------------------- +# Post-process passes (normalization stage), in end-state execution order. +# Faithful ports of the legacy __post_init__ handlers; each is invoked from +# its legacy slot via run_post_process_pass during the transition. +# --------------------------------------------------------------------------- + + +@register_post_process +def _sampling_backend_default(view: Any) -> dict: + if view.sampling_backend is None: + return { + "sampling_backend": ( + "flashinfer" if is_flashinfer_available() else "pytorch" + ) + } + return {} + + +@register_post_process +def _deterministic_sampling_backend(view: Any) -> dict: + if view.enable_deterministic_inference and view.sampling_backend != "ascend": + logger.warning( + "Sampling backend is set to pytorch for deterministic inference." + ) + return {"sampling_backend": "pytorch"} + return {} + + @dataclasses.dataclass(frozen=True) class OverrideRecord: """Provenance of one resolved write: ``base`` is the value before this diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 246f68fed..5c07ce1dd 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -314,6 +314,7 @@ class Flags(_StaticFlags): enable_multi_layer_eagle: bool = False swa_full_tokens_ratio: float = 0.8 disable_hybrid_swa_memory: bool = False + sampling_backend: str | None = None def freeze(self) -> None: for field in dataclasses.fields(self): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f6131faef..e05765080 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1430,6 +1430,7 @@ class ServerArgs: Arg( help="Choose the kernels for sampling layers.", choices=SAMPLING_BACKEND_CHOICES, + model_overridable=True, ), ] = None grammar_backend: A[ @@ -4611,10 +4612,14 @@ class ServerArgs: self._validate_mamba_no_buffer(model_arch) def _handle_sampling_backend(self): - if self.sampling_backend is None: - self.sampling_backend = ( - "flashinfer" if is_flashinfer_available() else "pytorch" - ) + # Moved to the resolution pipeline (arg_groups/overrides.py: + # _sampling_backend_default), invoked here at its legacy slot. + from sglang.srt.arg_groups.overrides import ( + _sampling_backend_default, + run_post_process_pass, + ) + + run_post_process_pass(self, _sampling_backend_default) def _get_default_attn_backend(self, use_mla_backend: bool, model_config): """ @@ -6307,12 +6312,15 @@ class ServerArgs: ) self.flashinfer_allreduce_fusion_backend = None - # Check sampling backend - if self.sampling_backend != "ascend": - self.sampling_backend = "pytorch" - logger.warning( - "Sampling backend is set to pytorch for deterministic inference." - ) + # The forced-pytorch sampling write moved to the resolution + # pipeline (arg_groups/overrides.py: + # _deterministic_sampling_backend), invoked at its legacy slot. + from sglang.srt.arg_groups.overrides import ( + _deterministic_sampling_backend, + run_post_process_pass, + ) + + run_post_process_pass(self, _deterministic_sampling_backend) is_deepseek_model = False if parse_connector_type(self.model_path) != ConnectorType.INSTANCE: try: diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index a4ae88219..f3908ce84 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -67,6 +67,7 @@ class TestModelOverridableWhitelist(CustomTestCase): "enable_multi_layer_eagle", "swa_full_tokens_ratio", "disable_hybrid_swa_memory", + "sampling_backend", } ), ) @@ -162,6 +163,53 @@ class TestModelOverrideRegistry(_IsolatedRegistry): ) +class TestResolvedViewAndPasses(CustomTestCase): + """Pipeline skeleton: read-only view semantics + transition invocation.""" + + def test_view_forwards_reads_and_rejects_writes(self): + from sglang.srt.arg_groups.overrides import ResolvedView + + live = SimpleNamespace(a=1, method=lambda: "m") + view = ResolvedView(live) + self.assertEqual(view.a, 1) + self.assertEqual(view.method(), "m") # method forwarding + live.a = 2 + self.assertEqual(view.a, 2) # live, not a snapshot + with self.assertRaises(AttributeError): + view.a = 3 + + def test_view_overlay_wins(self): + from sglang.srt.arg_groups.overrides import ResolvedView + + view = ResolvedView(SimpleNamespace(a=1, b=2), overlay={"a": 10}) + self.assertEqual(view.a, 10) + self.assertEqual(view.b, 2) + + def test_run_pass_appends_stash_and_dual_applies(self): + from sglang.srt.arg_groups.overrides import run_post_process_pass + + live = SimpleNamespace(x=None, _resolved_overrides=[]) + + def _fill_x(view): + return {"x": "filled"} if view.x is None else {} + + run_post_process_pass(live, _fill_x) + self.assertEqual(live.x, "filled") # dual-applied in place + self.assertEqual( + live._resolved_overrides, [(_fill_x.__qualname__, {"x": "filled"})] + ) + run_post_process_pass(live, _fill_x) # now a no-op + self.assertEqual(len(live._resolved_overrides), 1) + + def test_run_pass_rejects_non_dict(self): + from sglang.srt.arg_groups.overrides import run_post_process_pass + + with self.assertRaises(TypeError): + run_post_process_pass( + SimpleNamespace(_resolved_overrides=[]), lambda view: None + ) + + @dataclasses.dataclass class _FakeAttnGroup(_StaticFlags): backend: str = "unset" @@ -346,9 +394,9 @@ class TestGoldenModelOverrides(_IsolatedPublish): def test_mistral_large3_forces_bfloat16(self): sa = self._construct("MistralLarge3ForCausalLM", "mistral") self.assertEqual(sa.dtype, "bfloat16") # dual-apply == legacy write - self.assertEqual( + self.assertIn( + ("MODEL_OVERRIDES['MistralLarge3ForCausalLM']", {"dtype": "bfloat16"}), sa._resolved_overrides, - [("MODEL_OVERRIDES['MistralLarge3ForCausalLM']", {"dtype": "bfloat16"})], ) self.assertEqual(self._publish(sa).dtype, "bfloat16") @@ -368,7 +416,8 @@ class TestGoldenModelOverrides(_IsolatedPublish): def test_control_arch_keeps_pristine_dtype(self): sa = self._construct("LlamaForCausalLM", "llama") self.assertEqual(sa.dtype, "auto") - self.assertEqual(sa._resolved_overrides, []) + declared = {f for _s, d in sa._resolved_overrides for f in d} + self.assertNotIn("dtype", declared) # no arch declaration for Llama # publish still materializes the whitelisted leaf with the pristine # value: readers only ever read flags. self.assertEqual(self._publish(sa).dtype, "auto") @@ -376,9 +425,9 @@ class TestGoldenModelOverrides(_IsolatedPublish): def test_minimax_m2_enables_tf32_matmul(self): sa = self._construct("MiniMaxM2ForCausalLM", "llama") self.assertTrue(sa.enable_tf32_matmul) # dual-apply == legacy write - self.assertEqual( + self.assertIn( + ("_minimax_m2_overrides", {"enable_tf32_matmul": True}), sa._resolved_overrides, - [("_minimax_m2_overrides", {"enable_tf32_matmul": True})], ) flags = self._publish(sa) self.assertTrue(flags.enable_tf32_matmul) @@ -430,9 +479,9 @@ class TestGoldenModelOverrides(_IsolatedPublish): def test_gemma2_disables_hybrid_swa_memory(self): sa = self._construct("Gemma2ForCausalLM", "llama") self.assertTrue(sa.disable_hybrid_swa_memory) # dual-apply == legacy - self.assertEqual( + self.assertIn( + ("_gemma2_gemma3_overrides", {"disable_hybrid_swa_memory": True}), sa._resolved_overrides, - [("_gemma2_gemma3_overrides", {"disable_hybrid_swa_memory": True})], ) self.assertTrue(self._publish(sa).disable_hybrid_swa_memory) @@ -490,6 +539,46 @@ class TestGoldenModelOverrides(_IsolatedPublish): SimpleNamespace(architectures=["GptOssForCausalLM"]), ) + def test_sampling_backend_default_pass(self): + from sglang.srt.utils.common import is_flashinfer_available + + sa = self._construct("LlamaForCausalLM", "llama") + expected = "flashinfer" if is_flashinfer_available() else "pytorch" + self.assertEqual(sa.sampling_backend, expected) + self.assertIn( + ("_sampling_backend_default", {"sampling_backend": expected}), + sa._resolved_overrides, + ) + self.assertEqual(self._publish(sa).sampling_backend, expected) + + def test_sampling_backend_user_choice_survives(self): + sa = self._construct("LlamaForCausalLM", "llama", sampling_backend="pytorch") + self.assertEqual(sa.sampling_backend, "pytorch") + # the pass declared nothing; publish materializes the pristine choice + self.assertEqual(self._publish(sa).sampling_backend, "pytorch") + + def test_deterministic_inference_forces_pytorch_sampling(self): + sa = self._construct( + "LlamaForCausalLM", "llama", enable_deterministic_inference=True + ) + # two pass writers chain: default fill, then the deterministic force — + # last writer wins on the flags leaf and parity holds end-to-end. + self.assertEqual(sa.sampling_backend, "pytorch") + self.assertEqual(self._publish(sa).sampling_backend, "pytorch") + + def test_deterministic_ascend_is_left_alone(self): + from sglang.srt.arg_groups.overrides import ( + ResolvedView, + _deterministic_sampling_backend, + ) + + view = ResolvedView( + SimpleNamespace( + enable_deterministic_inference=True, sampling_backend="ascend" + ) + ) + self.assertEqual(_deterministic_sampling_backend(view), {}) + def test_step3p_declarations_at_callable_level(self): from sglang.srt.arg_groups.overrides import _step3p_overrides