From df55e911d6212e7ebdd66cb6738b88fdd1040ef2 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 1 Aug 2026 08:57:28 -0700 Subject: [PATCH] Fix the chunked-prefix-cache gate writing config the backends never read (#33168) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../sglang/srt/model_executor/model_runner.py | 11 ++- .../model_runner_components/misc_utils.py | 11 +-- .../runner/prefill_cuda_graph_runner.py | 5 +- .../srt/speculative/draft_worker_common.py | 36 ++++++---- .../test_chunked_prefix_cache_gate.py | 72 +++++++++++++++++++ .../unit/test_server_args_writer_ratchet.py | 2 +- 6 files changed, 109 insertions(+), 28 deletions(-) create mode 100644 test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 319ce8db4..388845d25 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -310,17 +310,16 @@ class ModelRunner: if server_args.show_time_cost: enable_show_time_cost() - misc_utils.maybe_disable_chunked_prefix_cache( - server_args=server_args, - use_mla_backend=self.use_mla_backend, - is_draft_worker=self.is_draft_worker, - ) - # Set the global server_args in the scheduler process (target worker # only, so a draft init cannot clobber target-derived global state). if not self.is_draft_worker: set_global_server_args_for_scheduler(server_args) + misc_utils.maybe_disable_chunked_prefix_cache( + use_mla_backend=self.use_mla_backend, + is_draft_worker=self.is_draft_worker, + ) + # Init OpenMP threads binding for CPU if self.device == "cpu": self.init_threads_binding() diff --git a/python/sglang/srt/model_executor/model_runner_components/misc_utils.py b/python/sglang/srt/model_executor/model_runner_components/misc_utils.py index f0261d3ae..d2308d123 100644 --- a/python/sglang/srt/model_executor/model_runner_components/misc_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/misc_utils.py @@ -4,6 +4,7 @@ import logging from typing import TYPE_CHECKING, Any, Optional from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa +from sglang.srt.runtime_context import get_context, get_exec, get_schedule from sglang.srt.server_args import CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS if TYPE_CHECKING: @@ -14,7 +15,7 @@ logger = logging.getLogger(__name__) def maybe_disable_chunked_prefix_cache( - *, server_args: ServerArgs, use_mla_backend: bool, is_draft_worker: bool + *, use_mla_backend: bool, is_draft_worker: bool ) -> None: # Chunked prefix caching requires an MLA model on a backend whose # kernels read that layout. This is a load-time gate, not a @@ -26,15 +27,15 @@ def maybe_disable_chunked_prefix_cache( return if ( not use_mla_backend - or server_args.attention_backend + or get_exec().kernel.attention_backend not in CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS ): - if not server_args.disable_chunked_prefix_cache: - server_args.override( + if not get_schedule().disable_chunked_prefix_cache: + get_context().override( "model_runner.chunked_prefix_cache_gate", disable_chunked_prefix_cache=True, ) - if not server_args.disable_chunked_prefix_cache: + if not get_schedule().disable_chunked_prefix_cache: logger.info("Chunked prefix cache is turned on.") diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 79b033f2b..5250dc1dd 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -103,7 +103,7 @@ from sglang.srt.model_executor.runner_utils.buffers import ( PrefillInputBuffers, ) from sglang.srt.model_loader.utils import resolve_language_model -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_schedule from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim from sglang.srt.utils import ( get_available_gpu_memory, @@ -403,8 +403,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # This flag controls whether the model dispatches through the distinct # chunked-prefix topology; backend capability is validated separately. self._capture_chunked_prefix = ( - self._is_full_backend - and not model_runner.server_args.disable_chunked_prefix_cache + self._is_full_backend and not get_schedule().disable_chunked_prefix_cache ) self._prefix_chunk_len = 0 self._prefix_chunk_capacity = 0 diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index 8dfb75d19..0aeb6e9e8 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -10,7 +10,7 @@ import torch from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode -from sglang.srt.runtime_context import get_context +from sglang.srt.runtime_context import get_context, get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 @@ -61,6 +61,26 @@ def _resolve_draft_attention_backend_fallback( return draft_backend +def draft_server_args_overrides(target_model_config, draft_backend) -> dict: + """Pre-publish field adjustments for a draft ``ServerArgs`` copy. + + Downstream draft-worker logic keys on ``speculative_draft_attention_backend`` + (backend selection in ``_get_attention_backend``, the fa4-draft KV dtype + override in ``configure_kv_cache_dtype``); ``context_length`` keeps the + draft aligned with the target; ``disable_chunked_prefix_cache`` is the + target's resolved gate (bag-only, absent from the pristine copy). + """ + return dict( + skip_tokenizer_init=True, + speculative_draft_attention_backend=draft_backend, + prefill_attention_backend=None, + decode_attention_backend=None, + attention_backend=draft_backend, + context_length=target_model_config.context_len, + disable_chunked_prefix_cache=get_schedule().disable_chunked_prefix_cache, + ) + + def build_draft_tp_worker( *, server_args: ServerArgs, @@ -81,20 +101,10 @@ def build_draft_tp_worker( ) ) # Post-resolution ServerArgs rejects bare assignment; route the draft-copy - # adjustments through the audited mutation point. Keep the resolved value - # on speculative_draft_attention_backend: downstream draft-worker logic - # keys on that field (backend selection in _get_attention_backend and the - # fa4-draft KV dtype override in configure_kv_cache_dtype), so nulling it - # would silently skip those paths. context_length keeps the draft aligned - # with the target. + # adjustments through the audited mutation point. draft_server_args.override( "draft_worker.build", - skip_tokenizer_init=True, - speculative_draft_attention_backend=draft_backend, - prefill_attention_backend=None, - decode_attention_backend=None, - attention_backend=draft_backend, - context_length=target_model_config.context_len, + **draft_server_args_overrides(target_model_config, draft_backend), ) # The draft's layers must resolve config from the draft's own bags. diff --git a/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py b/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py new file mode 100644 index 000000000..9f78f7b08 --- /dev/null +++ b/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py @@ -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() diff --git a/test/registered/unit/test_server_args_writer_ratchet.py b/test/registered/unit/test_server_args_writer_ratchet.py index a151bfd2e..426bca4f2 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 = 39 +_BASELINE = 38 class TestServerArgsWriterRatchet(CustomTestCase):