From ebb1c88d232e43a14cf9cc00ebedfe0351c13f8d Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sun, 2 Aug 2026 21:22:05 -0700 Subject: [PATCH] config: stop writing config onto the published ServerArgs at three sites (#33334) Each of these wrote a value after resolution so a later reader would find it on the instance. None of them needed the instance: one write was redundant, and the two that carry a value the resolved-config readback reports move to get_context().override, which the readback overlays. - The SM100 GDN prefill default was written onto ServerArgs and read back one line later by initialize_linear_attn_config. It is now the return value of flashinfer_gdn_prefill_default, threaded into initialize_linear_attn_config (an explicit --linear-attn-prefill-backend still wins) and recorded with get_context().override so /server_info reports the backend in effect. - The XGrammar fallback recorded grammar_backend="none" on the instance. No code reads the field after the factory reads it once, but get_internal_state reports the whole resolved config, so the fallback now lands there instead: the readback tells the truth and the seed keeps the requested backend. - UnifiedRadixCache.init_hicache re-applied the direct-IO layout fixup that __post_init__ already applies: init_hicache only runs when hierarchical cache is on, which is exactly when _handle_hicache normalizes page_first to page_first_direct (pinned by test_hicache_io_backend_and_mem_layout_ compatibility::direct_with_page_first). Three fixtures reached the fixup by building ServerArgs(model_path="dummy"), whose resolution is skipped, so they now declare the layout resolution would have produced. Writer ratchet 34 -> 31. --- .../srt/constrained/base_grammar_backend.py | 4 +- .../layers/attention/attention_registry.py | 13 ++- .../layers/attention/linear/gdn_backend.py | 21 ++--- .../srt/layers/attention/linear/utils.py | 6 +- .../srt/mem_cache/unified_radix_cache.py | 11 --- .../constrained/test_base_grammar_backend.py | 14 ++- .../test_gdn_prefill_backend_policy.py | 21 +---- .../attention/test_linear_attn_config.py | 92 +++++++++++++++++++ .../test_unified_radix_cache_unittest.py | 3 + .../unit/test_server_args_writer_ratchet.py | 2 +- 10 files changed, 136 insertions(+), 51 deletions(-) create mode 100644 test/registered/unit/layers/attention/test_linear_attn_config.py diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py index 585c0f07c..70a544724 100644 --- a/python/sglang/srt/constrained/base_grammar_backend.py +++ b/python/sglang/srt/constrained/base_grammar_backend.py @@ -22,7 +22,7 @@ from typing import Dict, List, NamedTuple, Optional, Tuple import torch from sglang.srt.parser.reasoning_parser import ReasoningParser -from sglang.srt.runtime_context import get_resources +from sglang.srt.runtime_context import get_context, get_resources from sglang.srt.server_args import ServerArgs logger = logging.getLogger(__name__) @@ -360,7 +360,7 @@ def create_grammar_backend( "Falling back to grammar_backend='none'. " "Structured outputs (JSON schema, regex, EBNF) will not be available." ) - server_args.override("grammar.import_fallback", grammar_backend="none") + get_context().override("grammar.import_fallback", grammar_backend="none") return None elif name == "llguidance": from sglang.srt.constrained.llguidance_backend import GuidanceBackend diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index bbc7b57f8..6c5400dc2 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -13,6 +13,7 @@ from sglang.srt.configs.linear_attn_model_registry import ( get_linear_attn_config, import_backend_class, ) +from sglang.srt.runtime_context import get_context from sglang.srt.utils import get_device_capability, is_hip, is_musa, is_npu _is_musa = is_musa() @@ -353,7 +354,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ) from sglang.srt.layers.attention.linear.gdn_backend import ( GDNAttnBackend, - maybe_set_default_flashinfer_gdn_prefill, + flashinfer_gdn_prefill_default, ) else: from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import ( @@ -367,9 +368,15 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ) check_environments() + prefill_default = None if hybrid_gdn_config(runner.model_config) is not None and not is_npu(): - maybe_set_default_flashinfer_gdn_prefill(runner) - initialize_linear_attn_config(runner.server_args) + prefill_default = flashinfer_gdn_prefill_default(runner) + if prefill_default is not None: + get_context().override( + "gdn_backend.sm100_flashinfer_default", + linear_attn_prefill_backend=prefill_default, + ) + initialize_linear_attn_config(runner.server_args, prefill_default) hybrid_backend_cls = HybridLinearAttnBackend if hybrid_gdn_config(runner.model_config) is not None: if is_blackwell(): diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 20af6d07f..d289dc2ee 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -63,8 +63,8 @@ elif is_cpu(): fused_gdn_gating = torch.ops.sgl_kernel.fused_gdn_gating_cpu -def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None: - """Use FlashInfer for the narrow SM100 GDN prefill domain we validated.""" +def flashinfer_gdn_prefill_default(model_runner: ModelRunner) -> Optional[str]: + """FlashInfer for the narrow SM100 GDN prefill domain we validated, else None.""" args = model_runner.server_args if ( args.linear_attn_prefill_backend is not None @@ -73,7 +73,7 @@ def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None: or not is_cuda() or torch.cuda.get_device_capability()[0] != 10 ): - return + return None cuda_version = torch.version.cuda chunk_size = args.chunked_prefill_size @@ -89,20 +89,17 @@ def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None: or model_runner.req_to_token_pool.mamba_pool.mamba_cache.temporal.dtype != torch.bfloat16 ): - return + return None from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( is_flashinfer_gdn_prefill_available, ) - if is_flashinfer_gdn_prefill_available(): - # server_args is resolved (read-only) by the time backends initialize; - # route this load-time default through the audited mutation entry. - args.override( - "gdn_backend.sm100_flashinfer_default", - linear_attn_prefill_backend="flashinfer", - ) - rank0_log("Defaulting SM100 GDN prefill backend to FlashInfer.") + if not is_flashinfer_gdn_prefill_available(): + return None + + rank0_log("Defaulting SM100 GDN prefill backend to FlashInfer.") + return "flashinfer" class GDNKernelDispatcher: diff --git a/python/sglang/srt/layers/attention/linear/utils.py b/python/sglang/srt/layers/attention/linear/utils.py index 2d7bb7858..520c10c78 100644 --- a/python/sglang/srt/layers/attention/linear/utils.py +++ b/python/sglang/srt/layers/attention/linear/utils.py @@ -43,13 +43,15 @@ LINEAR_ATTN_DECODE_BACKEND: Optional[LinearAttnKernelBackend] = None LINEAR_ATTN_PREFILL_BACKEND: Optional[LinearAttnKernelBackend] = None -def initialize_linear_attn_config(server_args: ServerArgs): +def initialize_linear_attn_config( + server_args: ServerArgs, prefill_default: Optional[str] = None +): global LINEAR_ATTN_DECODE_BACKEND global LINEAR_ATTN_PREFILL_BACKEND base = server_args.linear_attn_backend decode = server_args.linear_attn_decode_backend or base - prefill = server_args.linear_attn_prefill_backend or base + prefill = server_args.linear_attn_prefill_backend or prefill_default or base LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend(decode) LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend(prefill) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index de2405f89..c11a02e30 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -311,17 +311,6 @@ class UnifiedRadixCache(BasePrefixCache): attach_hybrid_pool_to_unified_cache, ) - # Direct IO layout fixup (must happen before pool creation) - if server_args.hicache_io_backend == "direct": - if server_args.hicache_mem_layout == "page_first": - server_args.override( - "hicache.mem_layout_force", hicache_mem_layout="page_first_direct" - ) - logger.warning( - "Page first layout is not supported with direct IO backend, " - "switching to page first direct layout" - ) - self.load_cache_event = threading.Event() self.sidecar_pool_specs.clear() self.extra_metric_labels = server_args.extra_metric_labels diff --git a/test/registered/unit/constrained/test_base_grammar_backend.py b/test/registered/unit/constrained/test_base_grammar_backend.py index 0cbce7ae0..13568c2cb 100644 --- a/test/registered/unit/constrained/test_base_grammar_backend.py +++ b/test/registered/unit/constrained/test_base_grammar_backend.py @@ -316,15 +316,21 @@ class TestCreateGrammarBackend(unittest.TestCase): @patch("sglang.srt.constrained.xgrammar_backend.XGrammarGrammarBackend") def test_xgrammar_unsupported_tokenizer_falls_back_to_none(self, mock_xgrammar_cls): from sglang.srt.constrained.xgrammar_backend import TokenizerNotSupportedError + from sglang.srt.runtime_context import get_context, get_exec mock_xgrammar_cls.side_effect = TokenizerNotSupportedError( "unsupported tokenizer" ) - args = self._make_server_args("xgrammar") + override = get_context().override_server_args(grammar_backend="xgrammar") + server_args = override.install() + self.addCleanup(override.restore) - result = create_grammar_backend(args, "tok", 32000, {1}) - self.assertIsNone(result) - self.assertEqual(args.grammar_backend, "none") + self.assertIsNone(create_grammar_backend(server_args, "tok", 32000, {1})) + self.assertEqual(get_exec().kernel.grammar_backend, "none") + self.assertEqual( + get_context().resolved_server_args_dict()["grammar_backend"], "none" + ) + self.assertEqual(server_args.grammar_backend, "xgrammar") @patch("sglang.srt.constrained.llguidance_backend.GuidanceBackend") def test_llguidance_backend(self, mock_guidance_cls): diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py index 95ea659e8..d10e0ae0f 100644 --- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py +++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py @@ -11,7 +11,7 @@ from sglang.srt.layers.attention.linear import gdn_backend from sglang.srt.layers.attention.linear.gdn_backend import ( GDNAttnBackend, GDNKernelDispatcher, - maybe_set_default_flashinfer_gdn_prefill, + flashinfer_gdn_prefill_default, ) from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( maybe_build_flashinfer_checkpoint_plan, @@ -39,11 +39,6 @@ def make_runner( enable_dynamic_chunking=False, chunked_prefill_size=8192, ) - # The policy routes its load-time default through the audited mutation entry - # (server_args.override); mirror that on the stub so the write lands. - args.override = MagicMock( - side_effect=lambda _source, **fields: vars(args).update(fields) - ) for name, value in arg_overrides.items(): setattr(args, name, value) @@ -87,16 +82,10 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): return_value=flashinfer_available, ), ): - maybe_set_default_flashinfer_gdn_prefill(runner) - return runner.server_args.linear_attn_prefill_backend + return flashinfer_gdn_prefill_default(runner) def test_selects_flashinfer_for_supported_sm100_gdn(self): - runner = make_runner() - self.assertEqual(self.apply_policy(runner), "flashinfer") - runner.server_args.override.assert_called_once_with( - "gdn_backend.sm100_flashinfer_default", - linear_attn_prefill_backend="flashinfer", - ) + self.assertEqual(self.apply_policy(make_runner()), "flashinfer") def test_selects_flashinfer_for_radix_cache_strategies(self): for strategy in ("no_buffer", "extra_buffer", "extra_buffer_lazy"): @@ -107,11 +96,11 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): ) self.assertEqual(self.apply_policy(runner), "flashinfer") - def test_preserves_explicit_prefill_override(self): + def test_declines_when_the_prefill_backend_is_explicit(self): for backend in ("triton", "flashinfer", "cutedsl"): with self.subTest(backend=backend): runner = make_runner(linear_attn_prefill_backend=backend) - self.assertEqual(self.apply_policy(runner), backend) + self.assertIsNone(self.apply_policy(runner)) def test_rejects_unsupported_capability(self): cases = ( diff --git a/test/registered/unit/layers/attention/test_linear_attn_config.py b/test/registered/unit/layers/attention/test_linear_attn_config.py new file mode 100644 index 000000000..41d78a788 --- /dev/null +++ b/test/registered/unit/layers/attention/test_linear_attn_config.py @@ -0,0 +1,92 @@ +"""Backend selection in initialize_linear_attn_config. + +The SM100 GDN default reaches the module state as an argument rather than as a +ServerArgs mutation, so the precedence between an explicit flag, that default, +and the shared base backend is pinned here. +""" + +import unittest + +from sglang.srt.layers.attention.linear import utils as linear_utils +from sglang.srt.layers.attention.linear.utils import ( + LinearAttnKernelBackend, + initialize_linear_attn_config, +) +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestLinearAttnConfig(CustomTestCase): + def setUp(self): + saved = ( + linear_utils.LINEAR_ATTN_DECODE_BACKEND, + linear_utils.LINEAR_ATTN_PREFILL_BACKEND, + ) + + def restore(): + ( + linear_utils.LINEAR_ATTN_DECODE_BACKEND, + linear_utils.LINEAR_ATTN_PREFILL_BACKEND, + ) = saved + + self.addCleanup(restore) + + def _init(self, prefill_default=None, **fields): + args = ServerArgs(model_path="dummy") + for key, value in fields.items(): + setattr(args, key, value) + initialize_linear_attn_config(args, prefill_default) + return ( + linear_utils.LINEAR_ATTN_PREFILL_BACKEND, + linear_utils.LINEAR_ATTN_DECODE_BACKEND, + ) + + def test_default_applies_when_the_flag_is_unset(self): + prefill, _ = self._init( + prefill_default="flashinfer", linear_attn_backend="triton" + ) + self.assertEqual(prefill, LinearAttnKernelBackend.FLASHINFER) + + def test_explicit_flag_wins_over_the_default(self): + prefill, _ = self._init( + prefill_default="flashinfer", + linear_attn_backend="triton", + linear_attn_prefill_backend="cutedsl", + ) + self.assertEqual(prefill, LinearAttnKernelBackend.CUTEDSL) + + def test_base_backend_applies_without_a_default(self): + prefill, decode = self._init(linear_attn_backend="triton") + self.assertEqual(prefill, LinearAttnKernelBackend.TRITON) + self.assertEqual(decode, LinearAttnKernelBackend.TRITON) + + def test_a_recorded_default_shows_in_the_resolved_config(self): + from sglang.srt.runtime_context import get_context, get_exec + + override = get_context().override_server_args(linear_attn_backend="triton") + server_args = override.install() + self.addCleanup(override.restore) + + get_context().override( + "gdn_backend.sm100_flashinfer_default", + linear_attn_prefill_backend="flashinfer", + ) + self.assertEqual(get_exec().mamba.linear_attn_prefill_backend, "flashinfer") + self.assertEqual( + get_context().resolved_server_args_dict()["linear_attn_prefill_backend"], + "flashinfer", + ) + self.assertIsNone(server_args.linear_attn_prefill_backend) + + def test_the_default_does_not_reach_the_decode_backend(self): + _, decode = self._init( + prefill_default="flashinfer", linear_attn_backend="triton" + ) + self.assertEqual(decode, LinearAttnKernelBackend.TRITON) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 8421dd38c..b875e4b9b 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -530,6 +530,7 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): model_path="dummy", page_size=self.cfg.page_size, hicache_io_backend="direct", + hicache_mem_layout="page_first_direct", hicache_write_policy=write_policy, ) set_global_server_args_for_scheduler(server_args) @@ -2869,6 +2870,7 @@ class UnifiedRadixCacheSuite: model_path="dummy", page_size=self.cfg.page_size, hicache_io_backend="direct", + hicache_mem_layout="page_first_direct", hicache_write_policy=write_policy, hicache_storage_backend=storage_backend, hicache_storage_backend_extra_config=storage_extra_config, @@ -6231,6 +6233,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): model_path="dummy", page_size=self.cfg.page_size, hicache_io_backend="direct", + hicache_mem_layout="page_first_direct", hicache_write_policy=write_policy, ) server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size) diff --git a/test/registered/unit/test_server_args_writer_ratchet.py b/test/registered/unit/test_server_args_writer_ratchet.py index b021f6b5c..3d2bde7e2 100644 --- a/test/registered/unit/test_server_args_writer_ratchet.py +++ b/test/registered/unit/test_server_args_writer_ratchet.py @@ -49,7 +49,7 @@ _EXCLUDED = ( "multimodal_gen", ) -_BASELINE = 34 +_BASELINE = 31 class TestServerArgsWriterRatchet(CustomTestCase):