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.
This commit is contained in:
Cheng Wan
2026-08-02 21:22:05 -07:00
committed by GitHub
parent 8186eeb939
commit ebb1c88d23
10 changed files with 136 additions and 51 deletions
@@ -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
@@ -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():
@@ -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:
@@ -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)
@@ -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
@@ -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):
@@ -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 = (
@@ -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()
@@ -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)
@@ -49,7 +49,7 @@ _EXCLUDED = (
"multimodal_gen",
)
_BASELINE = 34
_BASELINE = 31
class TestServerArgsWriterRatchet(CustomTestCase):