Fix the chunked-prefix-cache gate writing config the backends never read (#33168)
The load-time gate (maybe_disable_chunked_prefix_cache) wrote its ServerArgs instance while every reader has moved to the published config: the attention backends assert / branch on get_schedule().disable_chunked_prefix_cache when they initialize, so the flip never reached them and a backend outside CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS kept chunked prefix enabled. Reroute the writer through get_context().override (which writes the published bags) and flip the two remaining instance reads — the gate's own log check and the prefill cuda-graph runner's capture flag — to the bag. A regression test pins the three contracts: the gate lands on the bag, the pristine ServerArgs instance stays untouched, and the draft-worker guard never writes. The ServerArgs.override call-site ratchet drops 39 -> 38.
This commit is contained in:
@@ -0,0 +1,72 @@
|
||||
"""The load-time chunked-prefix gate must land on the published config bag.
|
||||
|
||||
Regression: the gate wrote the ServerArgs instance while every attention
|
||||
backend reads ``get_schedule().disable_chunked_prefix_cache`` — the bag never
|
||||
saw the flip, so an unsupported backend kept chunked prefix enabled.
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.srt.model_executor.model_runner_components.misc_utils import (
|
||||
maybe_disable_chunked_prefix_cache,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context, get_schedule, get_server_args
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestChunkedPrefixCacheGate(CustomTestCase):
|
||||
def _seed(self, **fields):
|
||||
override = get_context().override_server_args(**fields)
|
||||
override.install()
|
||||
self.addCleanup(override.restore)
|
||||
|
||||
def test_unsupported_backend_disables_on_the_bag(self):
|
||||
self._seed(attention_backend="triton")
|
||||
maybe_disable_chunked_prefix_cache(use_mla_backend=True, is_draft_worker=False)
|
||||
self.assertTrue(get_schedule().disable_chunked_prefix_cache)
|
||||
self.assertFalse(get_server_args().disable_chunked_prefix_cache)
|
||||
|
||||
def test_supported_backend_keeps_chunked_prefix(self):
|
||||
self._seed(attention_backend="fa3")
|
||||
maybe_disable_chunked_prefix_cache(use_mla_backend=True, is_draft_worker=False)
|
||||
self.assertFalse(get_schedule().disable_chunked_prefix_cache)
|
||||
|
||||
def test_draft_worker_never_writes(self):
|
||||
self._seed(attention_backend="triton")
|
||||
maybe_disable_chunked_prefix_cache(use_mla_backend=False, is_draft_worker=True)
|
||||
self.assertFalse(get_schedule().disable_chunked_prefix_cache)
|
||||
|
||||
def test_republish_discards_the_gate_so_it_must_run_after_publish(self):
|
||||
# Pins the ordering contract in ModelRunner.__init__: publishing
|
||||
# rebuilds the bags from the pristine instance, so the gate runs after
|
||||
# the target-worker publish.
|
||||
self._seed(attention_backend="triton")
|
||||
sa = get_server_args()
|
||||
maybe_disable_chunked_prefix_cache(use_mla_backend=True, is_draft_worker=False)
|
||||
self.assertTrue(get_schedule().disable_chunked_prefix_cache)
|
||||
get_context().set_server_args(sa) # what a later republish would do
|
||||
self.assertFalse(get_schedule().disable_chunked_prefix_cache)
|
||||
|
||||
def test_draft_copy_overrides_carry_the_gate(self):
|
||||
# The draft copy comes from the pristine instance, which never sees
|
||||
# the bag-only gate; the copy's pre-publish overrides carry it.
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.speculative.draft_worker_common import (
|
||||
draft_server_args_overrides,
|
||||
)
|
||||
|
||||
self._seed(attention_backend="triton")
|
||||
maybe_disable_chunked_prefix_cache(use_mla_backend=True, is_draft_worker=False)
|
||||
fields = draft_server_args_overrides(
|
||||
SimpleNamespace(context_len=64), draft_backend="fa3"
|
||||
)
|
||||
self.assertTrue(fields["disable_chunked_prefix_cache"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -49,7 +49,7 @@ _EXCLUDED = (
|
||||
"multimodal_gen",
|
||||
)
|
||||
|
||||
_BASELINE = 39
|
||||
_BASELINE = 38
|
||||
|
||||
|
||||
class TestServerArgsWriterRatchet(CustomTestCase):
|
||||
|
||||
Reference in New Issue
Block a user