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:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user