config: the alias form of the runner-side instance read
The previous batch counted `self.server_args.X` and called the runner surface done. It was not: the same read spelled through a local alias -- `server_args = model_runner.server_args` (or `sa = kvc.server_args`, `args = ...`) followed by `server_args.leaf` -- is the same process-global read wearing a local name, and the AST census counts **57 of them** across eleven files that the grep never saw. Census per function, following the alias. 52 were leaves and go to their bag (`spec` 11, `schedule` 9, `memory` 7, `exec.graph` 5, `exec.moe` 5, `parallel` 4, `disagg` 4, `model` 3, `exec.mamba` 2, `exec.overlap` 2). Five were not leaves: three derived members on the eager runner -- `max_speculative_num_draft_tokens` and `enable_mamba_extra_buffer` already had accessors, and `max_prefill_buffer_tokens` gets one (all its inputs are `schedule` leaves plus the configured PP size, so it derives from the bags and follows a post-publish override; `TestDerivedPredicatesAgreeAcrossTiers` pins it against the member over a 48-case matrix) -- plus `get_attention_backends()`, which the same commit routes through `attention_backends()`, and a dict that merely shares the name (`server_args_dict.items`). That dict is the one read left behind. `build_attention_backends` also stops resolving the pair from the record: it runs after publish, so it asks `attention_backends()` like every other consumer. The draft override on the runner still wins first. `dispatch_event_loop`'s three PP checks read the *configured* PP size, not the live topology: the MLX runner stub never initializes torch.distributed, so the live property asserts before the MLX event loop can start (a Codex catch). The configured leaf answers the same value wherever the live groups exist. `flashinfer_gdn_prefill_default`'s guard is the one read here that asks what the *operator* named rather than what the config resolved to, and the bag leaf now answers exactly that: the per-runner auto-default is stamped on the runner and deliberately never recorded process-wide, so nothing writes that leaf after launch and reading it back cannot mistake another runner's default for a flag. Three test doubles injected a `SimpleNamespace`/`MagicMock` record for exactly these reads and now publish instead (pool configurator, cache registry, GDN prefill policy) -- the fixture publishes what the case configures and hands the published instance to the whole-object contracts that still take one. The functions this sweep partially converted stop mixing sources (review catches): the flash-attention constructor's remaining seed reads (`speculative_eagle_topk`, `speculative_algorithm`, both deterministic gates) read their bags next to the leaves already converted; `_should_disable_scheduler_metadata_precompute` reads the parallel config leaves itself instead of taking the record (its alias binding was the last use); and the autotune gates (`disable_flashinfer_autotune`, deterministic, `flashinfer_autotune_skip_ops`) join the moe leaves the same function already reads from the bags. The pool-configurator fixture drops a parameter nothing published or read.
This commit is contained in:
Cheng Wan
@@ -24,14 +24,29 @@ from sglang.test.ci.ci_register import register_cpu_ci
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _publish(testcase, **fields):
|
||||
"""Install a published config for one case and restore on its cleanup --
|
||||
a failed case must not leave a partial publish for a later file in a
|
||||
monolithic local run."""
|
||||
from sglang.srt.runtime_context import get_context, get_server_args
|
||||
|
||||
override = get_context().override_server_args(**fields)
|
||||
override.install()
|
||||
testcase.addCleanup(override.restore)
|
||||
return get_server_args()
|
||||
|
||||
|
||||
def make_runner(
|
||||
testcase,
|
||||
*,
|
||||
state_dtype=torch.bfloat16,
|
||||
key_dim=128,
|
||||
value_dim=128,
|
||||
**arg_overrides,
|
||||
):
|
||||
args = SimpleNamespace(
|
||||
# The policy reads the published bags, so the fixture publishes the
|
||||
# configuration under test.
|
||||
fields = dict(
|
||||
linear_attn_backend="triton",
|
||||
linear_attn_prefill_backend=None,
|
||||
uses_mamba_radix_cache=False,
|
||||
@@ -40,8 +55,8 @@ def make_runner(
|
||||
enable_dynamic_chunking=False,
|
||||
chunked_prefill_size=8192,
|
||||
)
|
||||
for name, value in arg_overrides.items():
|
||||
setattr(args, name, value)
|
||||
fields.update(arg_overrides)
|
||||
args = _publish(testcase, **fields)
|
||||
|
||||
return SimpleNamespace(
|
||||
server_args=args,
|
||||
@@ -86,12 +101,13 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
|
||||
return flashinfer_gdn_prefill_default(runner)
|
||||
|
||||
def test_selects_flashinfer_for_supported_sm100_gdn(self):
|
||||
self.assertEqual(self.apply_policy(make_runner()), "flashinfer")
|
||||
self.assertEqual(self.apply_policy(make_runner(self)), "flashinfer")
|
||||
|
||||
def test_selects_flashinfer_for_radix_cache_strategies(self):
|
||||
for strategy in ("no_buffer", "extra_buffer", "extra_buffer_lazy"):
|
||||
with self.subTest(strategy=strategy):
|
||||
runner = make_runner(
|
||||
self,
|
||||
uses_mamba_radix_cache=True,
|
||||
mamba_radix_cache_strategy=strategy,
|
||||
)
|
||||
@@ -100,7 +116,7 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
|
||||
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)
|
||||
runner = make_runner(self, linear_attn_prefill_backend=backend)
|
||||
self.assertIsNone(self.apply_policy(runner))
|
||||
|
||||
def test_rejects_unsupported_capability(self):
|
||||
@@ -117,11 +133,11 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
|
||||
for name, runner_args, hardware in cases:
|
||||
with self.subTest(name=name):
|
||||
self.assertIsNone(
|
||||
self.apply_policy(make_runner(**runner_args), **hardware)
|
||||
self.apply_policy(make_runner(self, **runner_args), **hardware)
|
||||
)
|
||||
|
||||
def test_rejects_gdn_config_without_qwen_head_dims(self):
|
||||
runner = make_runner()
|
||||
runner = make_runner(self)
|
||||
runner.hybrid_gdn_config = SimpleNamespace()
|
||||
self.assertIsNone(self.apply_policy(runner))
|
||||
|
||||
@@ -136,7 +152,7 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
|
||||
)
|
||||
for name, runner_args in cases:
|
||||
with self.subTest(name=name):
|
||||
self.assertIsNone(self.apply_policy(make_runner(**runner_args)))
|
||||
self.assertIsNone(self.apply_policy(make_runner(self, **runner_args)))
|
||||
|
||||
def test_builds_compact_checkpoint_plan_for_packed_sequences(self):
|
||||
forward_batch = SimpleNamespace(
|
||||
|
||||
@@ -19,7 +19,18 @@ from sglang.srt.mem_cache.registry import (
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
def _publish(testcase, **fields):
|
||||
"""Install a published config for one case and restore on its cleanup."""
|
||||
from sglang.srt.runtime_context import get_context, get_server_args
|
||||
|
||||
override = get_context().override_server_args(**fields)
|
||||
override.install()
|
||||
testcase.addCleanup(override.restore)
|
||||
return get_server_args()
|
||||
|
||||
|
||||
def _make_ctx(
|
||||
testcase,
|
||||
*,
|
||||
backend=None,
|
||||
enable_streaming=False,
|
||||
@@ -32,11 +43,16 @@ def _make_ctx(
|
||||
effective_chunked_prefill_size=None,
|
||||
full_tokens_per_layer=None,
|
||||
):
|
||||
server_args = MagicMock()
|
||||
server_args.radix_cache_backend = backend
|
||||
server_args.enable_streaming_session = enable_streaming
|
||||
server_args.enable_lmcache = enable_lmcache
|
||||
server_args.enable_flexkv = False
|
||||
# The factory reads the published bags for the cache-backend leaves, so the
|
||||
# fixture publishes them; the instance stays for the whole-object contract
|
||||
# `TreeCacheBuildContext` carries.
|
||||
server_args = _publish(
|
||||
testcase,
|
||||
radix_cache_backend=backend,
|
||||
enable_streaming_session=enable_streaming,
|
||||
enable_lmcache=enable_lmcache,
|
||||
enable_flexkv=False,
|
||||
)
|
||||
return TreeCacheBuildContext(
|
||||
server_args=server_args,
|
||||
params=MagicMock(),
|
||||
@@ -101,14 +117,14 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase):
|
||||
factory = MagicMock(return_value=cache)
|
||||
register_radix_cache_backend("custom", factory)
|
||||
|
||||
result = create_tree_cache(_make_ctx(backend="custom"))
|
||||
result = create_tree_cache(_make_ctx(self, backend="custom"))
|
||||
|
||||
factory.assert_called_once()
|
||||
self.assertIs(result, cache)
|
||||
|
||||
def test_unknown_backend_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
create_tree_cache(_make_ctx(backend="not_a_real_backend"))
|
||||
create_tree_cache(_make_ctx(self, backend="not_a_real_backend"))
|
||||
|
||||
@patch("sglang.srt.mem_cache.registry.default_radix_cache_factory")
|
||||
def test_unset_backend_falls_back_to_default(self, default_factory):
|
||||
@@ -116,7 +132,7 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase):
|
||||
cache.supports_streaming_session.return_value = True
|
||||
default_factory.return_value = cache
|
||||
|
||||
result = create_tree_cache(_make_ctx(backend=None))
|
||||
result = create_tree_cache(_make_ctx(self, backend=None))
|
||||
|
||||
default_factory.assert_called_once()
|
||||
self.assertIs(result, cache)
|
||||
@@ -131,7 +147,7 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase):
|
||||
) as session_cls:
|
||||
session_cls.return_value = MagicMock(name="wrapped")
|
||||
result = create_tree_cache(
|
||||
_make_ctx(backend="nonstreaming", enable_streaming=True)
|
||||
_make_ctx(self, backend="nonstreaming", enable_streaming=True)
|
||||
)
|
||||
|
||||
session_cls.assert_called_once_with(inner)
|
||||
@@ -143,7 +159,7 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase):
|
||||
register_radix_cache_backend("streaming", MagicMock(return_value=inner))
|
||||
|
||||
result = create_tree_cache(
|
||||
_make_ctx(backend="streaming", enable_streaming=True)
|
||||
_make_ctx(self, backend="streaming", enable_streaming=True)
|
||||
)
|
||||
|
||||
self.assertIs(result, inner)
|
||||
@@ -158,7 +174,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
"""
|
||||
|
||||
def test_chunk_cache_when_chunked_prefill_and_disable_radix(self):
|
||||
ctx = _make_ctx(effective_chunked_prefill_size=512, disable_radix_cache=True)
|
||||
ctx = _make_ctx(
|
||||
self, effective_chunked_prefill_size=512, disable_radix_cache=True
|
||||
)
|
||||
with patch("sglang.srt.mem_cache.chunk_cache.ChunkCache") as ChunkCache:
|
||||
ChunkCache.return_value = MagicMock()
|
||||
result = default_radix_cache_factory(ctx)
|
||||
@@ -167,6 +185,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
|
||||
def test_swa_chunk_cache_when_chunked_prefill_disable_and_hybrid_swa(self):
|
||||
ctx = _make_ctx(
|
||||
self,
|
||||
effective_chunked_prefill_size=512,
|
||||
disable_radix_cache=True,
|
||||
is_hybrid_swa=True,
|
||||
@@ -179,6 +198,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
|
||||
def test_pure_swa_chunk_cache_when_chunked_prefill_disable_and_all_swa(self):
|
||||
ctx = _make_ctx(
|
||||
self,
|
||||
effective_chunked_prefill_size=512,
|
||||
disable_radix_cache=True,
|
||||
is_hybrid_swa=True,
|
||||
@@ -193,7 +213,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, PureSWAChunkCache.return_value)
|
||||
|
||||
def test_cpp_radix_cache_when_env_flag_set(self):
|
||||
ctx = _make_ctx()
|
||||
ctx = _make_ctx(
|
||||
self,
|
||||
)
|
||||
# `radix_cache_cpp` requires ninja + C++ extension to import, so
|
||||
# we inject a stand-in module rather than letting patch() trigger
|
||||
# the real import.
|
||||
@@ -215,7 +237,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_module.RadixCacheCpp.return_value)
|
||||
|
||||
def test_unified_radix_cache_when_env_flag_set(self):
|
||||
ctx = _make_ctx()
|
||||
ctx = _make_ctx(
|
||||
self,
|
||||
)
|
||||
# Shim both factory imports — each transitively loads sgl_kernel.
|
||||
fake_components = MagicMock()
|
||||
fake_radix = MagicMock()
|
||||
@@ -237,7 +261,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
|
||||
|
||||
def test_hi_radix_cache_when_hierarchical(self):
|
||||
ctx = _make_ctx(enable_hierarchical_cache=True)
|
||||
ctx = _make_ctx(self, enable_hierarchical_cache=True)
|
||||
# `hiradix_cache` imports `hicache_storage` and
|
||||
# `memory_pool_host`, both of which transitively load
|
||||
# `sgl_kernel`; inject a stand-in module.
|
||||
@@ -254,7 +278,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_module.HiRadixCache.return_value)
|
||||
|
||||
def test_unified_radix_cache_when_hierarchical_and_hybrid_ssm(self):
|
||||
ctx = _make_ctx(enable_hierarchical_cache=True, is_hybrid_ssm=True)
|
||||
ctx = _make_ctx(self, enable_hierarchical_cache=True, is_hybrid_ssm=True)
|
||||
# Hybrid SSM with hierarchical cache now uses UnifiedRadixCache.
|
||||
fake_components = MagicMock()
|
||||
fake_radix = MagicMock()
|
||||
@@ -274,7 +298,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
|
||||
|
||||
def test_unified_radix_cache_when_hierarchical_and_hybrid_swa(self):
|
||||
ctx = _make_ctx(enable_hierarchical_cache=True, is_hybrid_swa=True)
|
||||
ctx = _make_ctx(self, enable_hierarchical_cache=True, is_hybrid_swa=True)
|
||||
# Hybrid SWA with hierarchical cache also uses UnifiedRadixCache.
|
||||
fake_components = MagicMock()
|
||||
fake_radix = MagicMock()
|
||||
@@ -294,7 +318,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
|
||||
|
||||
def test_unified_radix_cache_when_hierarchical_and_dsa(self):
|
||||
ctx = _make_ctx(enable_hierarchical_cache=True, is_dsa=True)
|
||||
ctx = _make_ctx(self, enable_hierarchical_cache=True, is_dsa=True)
|
||||
# DSA models (e.g. DeepSeek V3.2 / GLM-5.1) with hierarchical cache
|
||||
# use UnifiedRadixCache.
|
||||
fake_components = MagicMock()
|
||||
@@ -315,7 +339,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
|
||||
|
||||
def test_swa_radix_cache_when_hybrid_swa(self):
|
||||
ctx = _make_ctx(is_hybrid_swa=True)
|
||||
ctx = _make_ctx(self, is_hybrid_swa=True)
|
||||
# SWA hybrid models now default to the unified radix tree.
|
||||
fake_components = MagicMock()
|
||||
fake_radix = MagicMock()
|
||||
@@ -331,7 +355,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
|
||||
|
||||
def test_pure_swa_radix_cache_when_all_swa(self):
|
||||
ctx = _make_ctx(is_hybrid_swa=True, full_tokens_per_layer=0)
|
||||
ctx = _make_ctx(self, is_hybrid_swa=True, full_tokens_per_layer=0)
|
||||
with patch(
|
||||
"sglang.srt.mem_cache.pure_swa_radix_cache.PureSWARadixCache"
|
||||
) as PureSWA:
|
||||
@@ -341,7 +365,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, PureSWA.return_value)
|
||||
|
||||
def test_mamba_radix_cache_when_hybrid_ssm(self):
|
||||
ctx = _make_ctx(is_hybrid_ssm=True)
|
||||
ctx = _make_ctx(self, is_hybrid_ssm=True)
|
||||
# Mamba hybrid models now default to the unified radix tree.
|
||||
fake_components = MagicMock()
|
||||
fake_radix = MagicMock()
|
||||
@@ -357,7 +381,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_radix.UnifiedRadixCache.return_value)
|
||||
|
||||
def test_lmc_radix_cache_when_enable_lmcache(self):
|
||||
ctx = _make_ctx(enable_lmcache=True)
|
||||
ctx = _make_ctx(self, enable_lmcache=True)
|
||||
# The lmcache backend raises at import time when the `lmcache`
|
||||
# package isn't installed, so inject a stand-in module instead
|
||||
# of letting patch() trigger the real import.
|
||||
@@ -377,7 +401,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
|
||||
self.assertIs(result, fake_module.LMCRadixCache.return_value)
|
||||
|
||||
def test_fallback_to_radix_cache(self):
|
||||
ctx = _make_ctx()
|
||||
ctx = _make_ctx(
|
||||
self,
|
||||
)
|
||||
with patch("sglang.srt.mem_cache.radix_cache.RadixCache") as RadixCache:
|
||||
RadixCache.return_value = MagicMock()
|
||||
result = default_radix_cache_factory(ctx)
|
||||
|
||||
@@ -11,7 +11,7 @@ from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.runtime_context import get_parallel, get_server_args
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
@@ -35,7 +35,21 @@ def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4):
|
||||
yield
|
||||
|
||||
|
||||
def _publish_config(testcase, **fields):
|
||||
"""Publish the configuration a runner double describes.
|
||||
|
||||
Installed per call and restored on the calling case's cleanup, so a failed
|
||||
case cannot leave a partial publish for a later file in a monolithic run.
|
||||
"""
|
||||
from sglang.srt.runtime_context import get_context
|
||||
|
||||
override = get_context().override_server_args(**fields)
|
||||
override.install()
|
||||
testcase.addCleanup(override.restore)
|
||||
|
||||
|
||||
def _make_model_runner(
|
||||
testcase,
|
||||
*,
|
||||
num_kv_heads=4,
|
||||
head_dim=64,
|
||||
@@ -56,7 +70,6 @@ def _make_model_runner(
|
||||
disable_overlap_schedule=False,
|
||||
sliding_window_size=None,
|
||||
speculative_num_draft_tokens=None,
|
||||
max_speculative_num_draft_tokens=None,
|
||||
speculative_algorithm=None,
|
||||
speculative_num_steps=None,
|
||||
speculative_eagle_topk=None,
|
||||
@@ -105,27 +118,29 @@ def _make_model_runner(
|
||||
mr.model_config = mc
|
||||
mr.kv_cache_dtype = "fake_bf16"
|
||||
|
||||
sa = SimpleNamespace()
|
||||
sa.max_total_tokens = None
|
||||
sa.swa_full_tokens_ratio = swa_full_tokens_ratio
|
||||
sa.page_size = page_size
|
||||
sa.disable_radix_cache = disable_radix_cache
|
||||
sa.chunked_prefill_size = chunked_prefill_size
|
||||
sa.disable_overlap_schedule = disable_overlap_schedule
|
||||
sa.speculative_num_draft_tokens = speculative_num_draft_tokens
|
||||
sa.max_speculative_num_draft_tokens = (
|
||||
max_speculative_num_draft_tokens or speculative_num_draft_tokens
|
||||
# The configurator reads the published bags, so the fixture publishes the
|
||||
# configuration it describes. The instance stays for the whole-object
|
||||
# hand-offs the configurator still does.
|
||||
_publish_config(
|
||||
testcase,
|
||||
max_total_tokens=None,
|
||||
swa_full_tokens_ratio=swa_full_tokens_ratio,
|
||||
page_size=page_size,
|
||||
disable_radix_cache=disable_radix_cache,
|
||||
chunked_prefill_size=chunked_prefill_size,
|
||||
disable_overlap_schedule=disable_overlap_schedule,
|
||||
speculative_num_draft_tokens=speculative_num_draft_tokens,
|
||||
speculative_algorithm=speculative_algorithm,
|
||||
speculative_num_steps=speculative_num_steps,
|
||||
speculative_eagle_topk=speculative_eagle_topk,
|
||||
disaggregation_mode=disaggregation_mode,
|
||||
max_running_requests=max_running_requests,
|
||||
disaggregation_decode_extra_slots=disaggregation_decode_extra_slots,
|
||||
enable_hisparse=False,
|
||||
enable_dsa_cache_layer_split=False,
|
||||
kv_cache_dtype="auto",
|
||||
)
|
||||
sa.speculative_algorithm = speculative_algorithm
|
||||
sa.speculative_num_steps = speculative_num_steps
|
||||
sa.speculative_eagle_topk = speculative_eagle_topk
|
||||
sa.disaggregation_mode = disaggregation_mode
|
||||
sa.max_running_requests = max_running_requests
|
||||
sa.disaggregation_decode_extra_slots = disaggregation_decode_extra_slots
|
||||
sa.enable_hisparse = False
|
||||
sa.enable_dsa_cache_layer_split = False
|
||||
sa.kv_cache_dtype = "auto"
|
||||
mr.server_args = sa
|
||||
mr.server_args = get_server_args()
|
||||
|
||||
spec = MagicMock()
|
||||
spec.is_eagle.return_value = False
|
||||
@@ -180,7 +195,7 @@ class TestDefaultConfigurator(unittest.TestCase):
|
||||
"""Default (MHA): available_bytes -> tokens, memory invariant holds."""
|
||||
|
||||
def _run(self, available_bytes, page_size=1, **kwargs):
|
||||
mr = _make_model_runner(page_size=page_size, **kwargs)
|
||||
mr = _make_model_runner(self, page_size=page_size, **kwargs)
|
||||
with mock_cpu_env():
|
||||
from sglang.srt.model_executor.pool_configurator import (
|
||||
create_memory_pool_configurator,
|
||||
@@ -238,10 +253,12 @@ class TestDefaultConfigurator(unittest.TestCase):
|
||||
):
|
||||
num_layers = 2
|
||||
raw = _make_model_runner(
|
||||
self,
|
||||
num_layers=num_layers,
|
||||
use_mla_backend=True,
|
||||
)
|
||||
packed = _make_model_runner(
|
||||
self,
|
||||
num_layers=num_layers,
|
||||
use_mla_backend=True,
|
||||
)
|
||||
@@ -265,6 +282,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
|
||||
|
||||
def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1):
|
||||
return _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=list(range(full_layers)),
|
||||
swa_attention_layer_ids=list(range(full_layers, full_layers + swa_layers)),
|
||||
@@ -353,6 +371,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
|
||||
def test_chunk_cache_cap_accounts_for_spec_topk_page_rounding(self):
|
||||
available = 1_000_000
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[0],
|
||||
swa_attention_layer_ids=[1],
|
||||
@@ -392,6 +411,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
|
||||
# 2*chunk(4) + page(1) = 9; cap = 41 * 2 + 9 = 91.
|
||||
available = 1_000_000
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[0],
|
||||
swa_attention_layer_ids=[1],
|
||||
@@ -422,6 +442,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
|
||||
def test_chunk_cache_cap_drops_prefill_for_disagg_decode(self):
|
||||
available = 1_000_000
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[0],
|
||||
swa_attention_layer_ids=[1],
|
||||
@@ -452,6 +473,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
|
||||
# under overlap.
|
||||
available = 1_000_000
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[0],
|
||||
swa_attention_layer_ids=[1],
|
||||
@@ -484,6 +506,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
|
||||
# request count (num_reserved_decode_tokens is a full-pool concern, not SWA).
|
||||
available = 2_000_000
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[0],
|
||||
swa_attention_layer_ids=[1],
|
||||
@@ -517,6 +540,7 @@ class TestAllSWAConfigurator(unittest.TestCase):
|
||||
|
||||
def _run(self, available_bytes, ratio=0.5, page_size=1, **kwargs):
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[],
|
||||
swa_attention_layer_ids=list(range(32)),
|
||||
@@ -569,7 +593,7 @@ class TestEagleConfigurator(unittest.TestCase):
|
||||
num_layers = 32
|
||||
eagle_draft_num_layers = 4
|
||||
|
||||
mr = _make_model_runner(num_layers=num_layers)
|
||||
mr = _make_model_runner(self, num_layers=num_layers)
|
||||
mr.spec_algorithm.is_eagle.return_value = True
|
||||
mr.spec_algorithm.is_standalone.return_value = False
|
||||
mr.spec_algorithm.is_none.return_value = False
|
||||
@@ -591,7 +615,7 @@ class TestEagleConfigurator(unittest.TestCase):
|
||||
|
||||
class TestFactory(unittest.TestCase):
|
||||
def test_default_for_non_swa(self):
|
||||
mr = _make_model_runner(is_hybrid_swa=False)
|
||||
mr = _make_model_runner(self, is_hybrid_swa=False)
|
||||
with mock_cpu_env():
|
||||
from sglang.srt.model_executor.pool_configurator import (
|
||||
DefaultPoolConfigurator,
|
||||
@@ -603,6 +627,7 @@ class TestFactory(unittest.TestCase):
|
||||
|
||||
def test_swa_for_hybrid(self):
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=list(range(16)),
|
||||
swa_attention_layer_ids=list(range(16, 32)),
|
||||
@@ -621,6 +646,7 @@ class TestFactory(unittest.TestCase):
|
||||
# SWAChunkCapPoolConfigurator is selected only when max_running_requests is set.
|
||||
def _cfg(max_running_requests):
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=[0],
|
||||
swa_attention_layer_ids=[1],
|
||||
@@ -681,7 +707,7 @@ class TestDflashDraftKvBudget(unittest.TestCase):
|
||||
def test_dcp_replication_scales_draft_budget(self):
|
||||
"""The replicated draft pool spans every DCP virtual location."""
|
||||
draft_kv_per_token = 10_240
|
||||
mr = _make_model_runner()
|
||||
mr = _make_model_runner(self)
|
||||
mr.spec_algorithm.is_dflash_family.return_value = True
|
||||
mr.spec_aux_config = SimpleNamespace(
|
||||
eagle_draft_num_layers=None,
|
||||
@@ -715,6 +741,7 @@ class TestDflashDraftKvBudget(unittest.TestCase):
|
||||
|
||||
def _tokens(draft_kv_per_token):
|
||||
mr = _make_model_runner(
|
||||
self,
|
||||
is_hybrid_swa=True,
|
||||
full_attention_layer_ids=list(range(16)),
|
||||
swa_attention_layer_ids=list(range(16, 32)),
|
||||
|
||||
@@ -57,6 +57,12 @@ _CONFIGURED_SIZE_CALL_SITES = {
|
||||
"point, since with PP off the group is never touched, which is what lets "
|
||||
"the Indexer be constructed before distributed init"
|
||||
),
|
||||
("srt/managers/scheduler.py", "configured_pp_size"): (
|
||||
"dispatch_event_loop picks the PP event loop; the MLX runner stub never "
|
||||
"initializes torch.distributed, so the live property asserts before the "
|
||||
"MLX loop can start -- the configured leaf answers the same value "
|
||||
"wherever the live groups exist"
|
||||
),
|
||||
("srt/mem_cache/kv_cache_configurator.py", "configured_pp_size"): (
|
||||
"decides whether the token capacity needs a cross-PP all-reduce at all; "
|
||||
"asking the configured size keeps that decision independent of whether a "
|
||||
|
||||
@@ -401,6 +401,7 @@ class _FakeResolvedArgs:
|
||||
max_running_requests: A[int | None, Arg(help="mrr"), NS("schedule")] = None
|
||||
chunked_prefill_size: A[int, Arg(help="cps"), NS("schedule")] = -1
|
||||
max_prefill_tokens: A[int, Arg(help="mpt"), NS("schedule")] = 16384
|
||||
enable_dynamic_chunking: A[bool, Arg(help="edc"), NS("schedule")] = False
|
||||
cuda_graph_config: A[object | None, Arg(help="cgc"), NS("exec.graph")] = None
|
||||
tp_size: A[int, Arg(help="tp"), NS("parallel")] = 1
|
||||
pp_size: A[int, Arg(help="pp"), NS("parallel")] = 1
|
||||
@@ -1026,6 +1027,31 @@ class TestDerivedPredicatesAgreeAcrossTiers(_IsolatedServerArgs):
|
||||
mamba_extra_buffer_lazy_enabled(),
|
||||
)
|
||||
|
||||
def test_prefill_buffer_ceiling_matches_the_member(self):
|
||||
from sglang.srt.runtime_context import max_prefill_buffer_tokens
|
||||
|
||||
for chunked in (-1, 0, 1024, 8192):
|
||||
for dynamic in (False, True):
|
||||
for pp in (1, 4):
|
||||
for max_prefill in (0, 2048, 16384):
|
||||
with self.subTest(
|
||||
chunked=chunked,
|
||||
dynamic=dynamic,
|
||||
pp=pp,
|
||||
max_prefill=max_prefill,
|
||||
):
|
||||
args = _FakeResolvedArgs(
|
||||
chunked_prefill_size=chunked,
|
||||
enable_dynamic_chunking=dynamic,
|
||||
pp_size=pp,
|
||||
max_prefill_tokens=max_prefill,
|
||||
)
|
||||
get_context().set_server_args(args)
|
||||
self.assertEqual(
|
||||
ServerArgs.max_prefill_buffer_tokens(args),
|
||||
max_prefill_buffer_tokens(),
|
||||
)
|
||||
|
||||
def test_activation_reserve_matches_the_member(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
|
||||
Reference in New Issue
Block a user