[refactor] Add the post-process resolution stage; migrate sampling_backend (stack 10/15) (#30072)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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

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):

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:

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