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:
Cheng Wan
2026-08-01 08:57:28 -07:00
committed by GitHub
parent ae84811666
commit df55e911d6
6 changed files with 109 additions and 28 deletions
@@ -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()
@@ -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.")
@@ -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
@@ -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.