From 110bf7e6a83b1b2bea74ffe52b7f6aa20f34b882 Mon Sep 17 00:00:00 2001
From: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Date: Sun, 9 Aug 2026 14:43:28 -0700
Subject: [PATCH] config: retire the hidden global fallbacks and the
 mamba-extra-buffer instance reads (#34080)

---
 python/sglang/srt/arg_groups/overrides.py     | 14 +++++++-
 python/sglang/srt/managers/schedule_batch.py  | 21 ++++++-----
 .../batch_result_processor.py                 |  5 +--
 .../sglang/srt/mem_cache/allocation_sizing.py | 11 +++---
 python/sglang/srt/models/deepseek_v2.py       | 11 ++----
 python/sglang/srt/models/inkling.py           |  5 ++-
 .../sglang/srt/models/inkling_common/sconv.py | 11 +++---
 python/sglang/srt/runtime_context.py          | 23 ++++++++++++
 python/sglang/srt/speculative/spec_utils.py   | 25 +++++++------
 .../mlx/test_attention_patching.py            | 12 ++++---
 ...t_batch_result_processor_mamba_boundary.py | 36 ++++++-------------
 .../test_schedule_batch_prepare_for_decode.py | 21 ++++++-----
 12 files changed, 112 insertions(+), 83 deletions(-)

diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py
index 36df0d9c1..7c009cf75 100644
--- a/python/sglang/srt/arg_groups/overrides.py
+++ b/python/sglang/srt/arg_groups/overrides.py
@@ -290,13 +290,25 @@ def attention_backends_of(cfg: Any) -> tuple:
 
 def mamba_extra_buffer_of(cfg: Any) -> bool:
     """Mid-resolution equivalent of runtime_context.mamba_extra_buffer_enabled:
-    reads the (possibly overlaid) strategy from a config-shaped object."""
+    reads the (possibly overlaid) strategy from a config-shaped object.
+
+    This is the one definition of the predicate: ``ServerArgs`` delegates its
+    member to it, and the runtime_context accessor is its post-publish sibling
+    (which cannot reuse it, because the two leaves land in different bags)."""
     return cfg.disable_radix_cache is False and cfg.mamba_radix_cache_strategy in (
         "extra_buffer",
         "extra_buffer_lazy",
     )
 
 
+def mamba_extra_buffer_lazy_of(cfg: Any) -> bool:
+    """The lazy variant of :func:`mamba_extra_buffer_of`."""
+    return (
+        cfg.disable_radix_cache is False
+        and cfg.mamba_radix_cache_strategy == "extra_buffer_lazy"
+    )
+
+
 def collect_model_override_declarations(
     architecture: str, server_args: Any, hf_config: Any
 ) -> List[Tuple[str, Dict[str, Any]]]:
diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py
index 0ae504235..78a534bff 100755
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -2,7 +2,14 @@ from __future__ import annotations
 
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-from sglang.srt.runtime_context import get_exec, get_schedule, get_serving, get_spec
+from sglang.srt.runtime_context import (
+    get_exec,
+    get_schedule,
+    get_serving,
+    get_spec,
+    mamba_extra_buffer_enabled,
+    mamba_extra_buffer_lazy_enabled,
+)
 from sglang.srt.utils.common import (
     Range,
     ceil_align,
@@ -2342,7 +2349,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     def prepare_for_extend(self):
         self.forward_mode = ForwardMode.EXTEND
-        server_args = get_server_args()
 
         if self.is_dllm():
             # For DLLM, we use a separate forward mode
@@ -2471,7 +2477,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 req.already_computed = seq_len
             req.is_retracted = False
 
-            if server_args.enable_mamba_extra_buffer():
+            if mamba_extra_buffer_enabled():
                 track_entry = self._mamba_radix_cache_v2_req_prepare_for_extend(req)
                 mamba_track_mask_cpu.append(track_entry.track_mask)
                 mamba_track_indices_cpu.append(track_entry.track_index)
@@ -2576,7 +2582,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         self.extend_logprob_start_lens = extend_logprob_start_lens
         self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
 
-        if server_args.enable_mamba_extra_buffer():
+        if mamba_extra_buffer_enabled():
             self.mamba_track_indices = torch.tensor(
                 mamba_track_indices_cpu,
                 dtype=torch.int64,
@@ -2662,7 +2668,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             # In lazy mode, skip the swap — the second ping-pong slot is not
             # allocated yet; it will be allocated on demand at the track boundary
             # in mamba_lazy_prealloc_at_boundary during prepare_for_decode.
-            if not server_args.enable_mamba_extra_buffer_lazy():
+            if not mamba_extra_buffer_lazy_enabled():
                 req.mamba_next_track_idx = (
                     self.req_to_token_pool.get_mamba_ping_pong_other_idx(
                         req.mamba_next_track_idx
@@ -3003,7 +3009,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     def prepare_for_decode(self):
         self.forward_mode = ForwardMode.DECODE
-        server_args = get_server_args()
         # Decode embeds the last output token via embed_tokens; clear the stale
         # prefill-time tensor so it doesn't leak into ForwardBatch.
         self.input_embeds = None
@@ -3057,7 +3062,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 self.req_pool_indices_cpu,
             )
 
-        if server_args.enable_mamba_extra_buffer():
+        if mamba_extra_buffer_enabled():
             mamba_track_interval = get_exec().mamba.mamba_track_interval
 
             if len(self.reqs) == 0:
@@ -3065,7 +3070,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                     (0,), dtype=torch.int64, device=self.device
                 )
             else:
-                if server_args.enable_mamba_extra_buffer_lazy():
+                if mamba_extra_buffer_lazy_enabled():
                     self.mamba_lazy_prealloc_at_boundary(mamba_track_interval)
                 set_mamba_track_indices_from_reqs(self)
 
diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py
index 0f4ce352f..73998604e 100644
--- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py
@@ -38,6 +38,7 @@ from sglang.srt.runtime_context import (
     get_memory,
     get_observability,
     get_server_args,
+    mamba_extra_buffer_lazy_enabled,
 )
 from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
 from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
@@ -1073,7 +1074,7 @@ class SchedulerBatchResultProcessor:
                     prepare_release(req)
                 is_insert = (
                     req.mamba_lazy_is_insert
-                    if get_server_args().enable_mamba_extra_buffer_lazy()
+                    if mamba_extra_buffer_lazy_enabled()
                     else True
                 )
                 release_kv_cache(req, self.tree_cache, is_insert=is_insert)
@@ -1109,7 +1110,7 @@ class SchedulerBatchResultProcessor:
         if req.mamba_ping_pong_track_buffer is None:
             return
 
-        lazy = get_server_args().enable_mamba_extra_buffer_lazy()
+        lazy = mamba_extra_buffer_lazy_enabled()
         if known_boundary:
             self._mamba_assert_committed_len_lookahead(req)
             track_seqlen = req.kv_committed_len
diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py
index 10d08d628..372c1fd87 100644
--- a/python/sglang/srt/mem_cache/allocation_sizing.py
+++ b/python/sglang/srt/mem_cache/allocation_sizing.py
@@ -6,10 +6,7 @@ from sglang.srt.runtime_context import get_server_args
 from sglang.srt.server_args import ServerArgs
 
 
-def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
-    if server_args is None:
-        server_args = get_server_args()
-
+def get_alloc_len_per_decode(server_args: ServerArgs) -> int:
     if server_args.speculative_algorithm is None:
         return 1
 
@@ -39,7 +36,13 @@ def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> in
 
     The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
     mode; see eagle_utils.eagle_prepare_for_decode.
+
+    Callers on a request path have no config in hand, so this is the module's
+    single "which config" decision point: everything below it takes the
+    instance explicitly.
     """
+    if server_args is None:
+        server_args = get_server_args()
     return 2 * get_alloc_len_per_decode(server_args)
 
 
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index 6fcbc1cb7..8c424e242 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -858,11 +858,7 @@ class DeepseekV2MoE(nn.Module):
             )
         ]
 
-    def _can_dual_stream_graph(
-        self, hidden_states: torch.Tensor, server_args=None
-    ) -> bool:
-        if server_args is None:
-            server_args = get_server_args()
+    def _can_dual_stream_graph(self, hidden_states: torch.Tensor) -> bool:
         return (
             _enable_pcg_dsv2_dual_stream
             and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
@@ -875,7 +871,7 @@ class DeepseekV2MoE(nn.Module):
             and not self._enable_a2a_moe
             and not self._fuse_shared_experts_inside_sbo
             and not getattr(self, "is_hash", False)
-            and not server_args.enable_eplb
+            and not get_exec().moe.enable_eplb
         )
 
     def forward(
@@ -898,8 +894,7 @@ class DeepseekV2MoE(nn.Module):
             )
 
         if not self._enable_a2a_moe:
-            server_args = get_server_args()
-            if self._can_dual_stream_graph(hidden_states, server_args):
+            if self._can_dual_stream_graph(hidden_states):
                 fwd = get_forward()
                 return dsv2_flashinfer_moe_dual_stream_graph(
                     hidden_states,
diff --git a/python/sglang/srt/models/inkling.py b/python/sglang/srt/models/inkling.py
index f3f5a610c..ae0d280cd 100644
--- a/python/sglang/srt/models/inkling.py
+++ b/python/sglang/srt/models/inkling.py
@@ -80,7 +80,7 @@ from sglang.srt.runtime_context import (
     get_model,
     get_parallel,
     get_schedule,
-    get_server_args,
+    mamba_extra_buffer_enabled,
 )
 from sglang.srt.utils import add_prefix, is_cuda, make_layers
 
@@ -1017,12 +1017,11 @@ class InklingForConditionalGeneration(nn.Module):
         self.config = config
         self.text_config = config.text_config
 
-        server_args = get_server_args()
         assert envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get()
         if get_disagg().disaggregation_mode != "decode":
             assert not get_memory().disable_radix_cache
             assert not get_schedule().disable_hybrid_swa_memory
-            assert server_args.enable_mamba_extra_buffer()
+            assert mamba_extra_buffer_enabled()
 
         from types import SimpleNamespace
 
diff --git a/python/sglang/srt/models/inkling_common/sconv.py b/python/sglang/srt/models/inkling_common/sconv.py
index 167fc7dad..a846da550 100644
--- a/python/sglang/srt/models/inkling_common/sconv.py
+++ b/python/sglang/srt/models/inkling_common/sconv.py
@@ -17,7 +17,11 @@ from sglang.srt.models.inkling_common.kernels.sconv import (
     save_intermediate_conv_windows,
     update_sconv_cache,
 )
-from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args
+from sglang.srt.runtime_context import (
+    get_exec,
+    get_parallel,
+    mamba_extra_buffer_enabled,
+)
 from sglang.srt.utils import is_cuda, set_weight_attrs
 
 
@@ -267,10 +271,7 @@ class ShortConvolution(nn.Module):
             draft_token_num = hidden_states.shape[1]
 
         mamba_track_indices = getattr(forward_batch, "mamba_track_indices", None)
-        do_tracking = (
-            mamba_track_indices is not None
-            and get_server_args().enable_mamba_extra_buffer()
-        )
+        do_tracking = mamba_track_indices is not None and mamba_extra_buffer_enabled()
 
         crossed = track_step = None
         if do_tracking:
diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py
index b501d91a4..c780ddb61 100644
--- a/python/sglang/srt/runtime_context.py
+++ b/python/sglang/srt/runtime_context.py
@@ -1383,3 +1383,26 @@ def reset_context() -> None:
     _CONTEXT.resources = Resources()
     _CONTEXT.forward = ForwardFlags()
     set_global_dwdp_manager(None)
+
+
+def mamba_extra_buffer_enabled() -> bool:
+    """Whether the mamba radix cache keeps its extra state buffer.
+
+    A predicate over two published leaves (``memory.disable_radix_cache`` and
+    ``exec.mamba.mamba_radix_cache_strategy``), so it reads the bags rather
+    than the startup record — the ``ServerArgs`` member of the same name is the
+    pre-publish equivalent used inside the resolution pipeline.
+    """
+    return (
+        get_memory().disable_radix_cache is False
+        and get_exec().mamba.mamba_radix_cache_strategy
+        in ("extra_buffer", "extra_buffer_lazy")
+    )
+
+
+def mamba_extra_buffer_lazy_enabled() -> bool:
+    """The lazy variant of :func:`mamba_extra_buffer_enabled`."""
+    return (
+        get_memory().disable_radix_cache is False
+        and get_exec().mamba.mamba_radix_cache_strategy == "extra_buffer_lazy"
+    )
diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py
index bb5aeb94a..1c9044292 100644
--- a/python/sglang/srt/speculative/spec_utils.py
+++ b/python/sglang/srt/speculative/spec_utils.py
@@ -46,7 +46,13 @@ from sglang.srt.mem_cache.allocation import (
 from sglang.srt.mem_cache.allocation import (
     assign_req_to_token_pool_func as assign_req_to_token_pool_func,
 )
-from sglang.srt.runtime_context import get_exec, get_server_args
+from sglang.srt.runtime_context import (
+    get_exec,
+    get_server_args,
+    get_spec,
+    mamba_extra_buffer_enabled,
+    mamba_extra_buffer_lazy_enabled,
+)
 from sglang.srt.utils import (
     is_cpu,
     is_cuda,
@@ -250,16 +256,14 @@ def record_stream_for_v2_verify(batch, verify_input, fwd_stream):
     record_stream_each(candidates, fwd_stream)
 
 
-def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
-    if server_args is None:
-        server_args = get_server_args()
-
+def spec_need_hidden_states() -> bool:
     # STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
     # multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap.
     # TODO(lsyin): also skip when step == 1.
-    if server_args.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"):
+    spec = get_spec()
+    if spec.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"):
         return False
-    return not server_args.enable_multi_layer_eagle
+    return not spec.enable_multi_layer_eagle
 
 
 @torch.compile(dynamic=True, disable=_is_npu or _is_xpu)
@@ -761,11 +765,10 @@ def prepare_mamba_track_for_verify(batch: ScheduleBatch) -> None:
     Lazy: gather the positions planned by mamba_lazy_spec_prepare. Runs
     inside forward isolation, so it must not mutate req/pool state.
     """
-    server_args = get_server_args()
-    if not server_args.enable_mamba_extra_buffer():
+    if not mamba_extra_buffer_enabled():
         return
     track_positions = None
-    if server_args.enable_mamba_extra_buffer_lazy():
+    if mamba_extra_buffer_lazy_enabled():
         track_positions = batch.mamba_lazy_spec_track_positions_cpu
         assert track_positions is not None and len(track_positions) == len(
             batch.reqs
@@ -1028,7 +1031,7 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
     prep on its draft input -- the dispatcher routes.
     """
     server_args = get_server_args()
-    if server_args.enable_mamba_extra_buffer_lazy():
+    if mamba_extra_buffer_lazy_enabled():
         # Scheduler phase (outside forward isolation).
         batch.mamba_lazy_spec_prepare(
             get_exec().mamba.mamba_track_interval,
diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py
index 8ce22bedd..921eeb541 100644
--- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py
+++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py
@@ -1232,7 +1232,6 @@ class TestMlxOverlapScheduler(unittest.TestCase):
         logits_output = SimpleNamespace(customized_info=None)
         original_release = batch_result_processor_module.release_kv_cache
         original_get_indexer = batch_result_processor_module.get_global_indexer_capturer
-        original_get_server_args = batch_result_processor_module.get_server_args
 
         def fake_release_kv_cache(release_req, tree_cache, is_insert=False):
             events.append(("release", release_req.rid))
@@ -1240,21 +1239,26 @@ class TestMlxOverlapScheduler(unittest.TestCase):
 
         batch_result_processor_module.release_kv_cache = fake_release_kv_cache
         batch_result_processor_module.get_global_indexer_capturer = lambda: None
-        batch_result_processor_module.get_server_args = lambda: SimpleNamespace(
-            enable_mamba_extra_buffer_lazy=lambda: False
+        # The lazy predicate reads the published bags; publish the non-lazy
+        # strategy instead of stubbing the accessor.
+        from sglang.srt.runtime_context import get_context
+
+        override = get_context().override_server_args(
+            mamba_radix_cache_strategy="extra_buffer"
         )
+        override.install()
         try:
             SchedulerBatchResultProcessor._handle_finish_state_updated_req(
                 processor, req, batch, result, i, logits_output
             )
         finally:
+            override.restore()
             for name, original in saved.items():
                 setattr(SchedulerBatchResultProcessor, name, original)
             batch_result_processor_module.release_kv_cache = original_release
             batch_result_processor_module.get_global_indexer_capturer = (
                 original_get_indexer
             )
-            batch_result_processor_module.get_server_args = original_get_server_args
 
         self.assertEqual(
             events,
diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py
index a40308602..17431c6c9 100644
--- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py
+++ b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py
@@ -10,6 +10,7 @@ from sglang.srt.managers.scheduler import Scheduler
 from sglang.srt.managers.scheduler_components.batch_result_processor import (
     SchedulerBatchResultProcessor,
 )
+from sglang.srt.runtime_context import get_context
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -144,42 +145,25 @@ class TestMambaBoundaryMaskReuse(unittest.TestCase):
                     processor.process_batch_result_decode(result_batch, batch_result)
 
                 scheduler.process_batch_result = process_batch_result
-                server_args = SimpleNamespace(
-                    enable_mamba_extra_buffer=lambda: True,
-                    enable_mamba_extra_buffer_lazy=lambda: False,
-                )
 
                 with (
+                    # The mamba predicates and the track interval read the
+                    # published bags, so publish the configuration under test
+                    # (non-lazy extra buffer, interval 4); observability and
+                    # disagg reads are served by the same publish at their
+                    # defaults.
+                    get_context().override_server_args(
+                        mamba_radix_cache_strategy="extra_buffer",
+                        mamba_track_interval=4,
+                    ),
                     patch(
                         "sglang.srt.managers.schedule_batch.alloc_for_decode",
                         return_value=torch.tensor([3], dtype=torch.int64),
                     ),
-                    patch(
-                        "sglang.srt.managers.schedule_batch.get_server_args",
-                        return_value=server_args,
-                    ),
-                    patch(
-                        "sglang.srt.managers.schedule_batch.get_exec",
-                        return_value=SimpleNamespace(
-                            mamba=SimpleNamespace(mamba_track_interval=4)
-                        ),
-                    ),
                     patch(
                         "sglang.srt.managers.schedule_batch.set_mamba_track_indices_from_reqs"
                     ),
                     patch.object(torch.Tensor, "pin_memory", lambda tensor: tensor),
-                    patch(
-                        "sglang.srt.managers.scheduler_components."
-                        "batch_result_processor.get_observability",
-                        return_value=SimpleNamespace(enable_metrics=False),
-                    ),
-                    patch(
-                        "sglang.srt.managers.scheduler_components."
-                        "batch_result_processor.get_disagg",
-                        return_value=SimpleNamespace(
-                            disaggregation_decode_enable_offload_kvcache=False
-                        ),
-                    ),
                     patch.object(
                         SchedulerBatchResultProcessor,
                         "_mamba_prefix_cache_update",
diff --git a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py
index a9b66b2bc..da5c6b489 100644
--- a/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py
+++ b/test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py
@@ -4,6 +4,7 @@ from unittest.mock import patch
 
 import torch
 
+from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import maybe_stub_sgl_kernel
 
@@ -43,18 +44,16 @@ class TestPrepareForDecodeSeqLensOwnership(unittest.TestCase):
         """Each prepare_for_decode call rebinds seq-lens tensors to new +1 objects without mutating the old ones."""
         batch = _make_decode_batch()
 
-        server_args = types.SimpleNamespace(
-            enable_mamba_extra_buffer=lambda: False,
+        # The mamba-extra-buffer predicate reads the published bags, so the
+        # fixture publishes a config with the strategy off.
+        override = get_context().override_server_args(
+            mamba_radix_cache_strategy="no_buffer"
         )
-        with (
-            patch(
-                "sglang.srt.managers.schedule_batch.alloc_for_decode",
-                return_value=torch.tensor([6, 7], dtype=torch.int64),
-            ),
-            patch(
-                "sglang.srt.managers.schedule_batch.get_server_args",
-                return_value=server_args,
-            ),
+        override.install()
+        self.addCleanup(override.restore)
+        with patch(
+            "sglang.srt.managers.schedule_batch.alloc_for_decode",
+            return_value=torch.tensor([6, 7], dtype=torch.int64),
         ):
             for step in range(1, 3):
                 prev_seq_lens = batch.seq_lens