config: retire the hidden global fallbacks and the mamba-extra-buffer instance reads (#34080)

This commit is contained in:
Cheng Wan
2026-08-09 14:43:28 -07:00
committed by GitHub
parent 967cac801c
commit 110bf7e6a8
12 changed files with 112 additions and 83 deletions
+13 -1
View File
@@ -290,13 +290,25 @@ def attention_backends_of(cfg: Any) -> tuple:
def mamba_extra_buffer_of(cfg: Any) -> bool:
"""Mid-resolution equivalent of runtime_context.mamba_extra_buffer_enabled:
reads the (possibly overlaid) strategy from a config-shaped object."""
reads the (possibly overlaid) strategy from a config-shaped object.
This is the one definition of the predicate: ``ServerArgs`` delegates its
member to it, and the runtime_context accessor is its post-publish sibling
(which cannot reuse it, because the two leaves land in different bags)."""
return cfg.disable_radix_cache is False and cfg.mamba_radix_cache_strategy in (
"extra_buffer",
"extra_buffer_lazy",
)
def mamba_extra_buffer_lazy_of(cfg: Any) -> bool:
"""The lazy variant of :func:`mamba_extra_buffer_of`."""
return (
cfg.disable_radix_cache is False
and cfg.mamba_radix_cache_strategy == "extra_buffer_lazy"
)
def collect_model_override_declarations(
architecture: str, server_args: Any, hf_config: Any
) -> List[Tuple[str, Dict[str, Any]]]:
+13 -8
View File
@@ -2,7 +2,14 @@ from __future__ import annotations
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.runtime_context import get_exec, get_schedule, get_serving, get_spec
from sglang.srt.runtime_context import (
get_exec,
get_schedule,
get_serving,
get_spec,
mamba_extra_buffer_enabled,
mamba_extra_buffer_lazy_enabled,
)
from sglang.srt.utils.common import (
Range,
ceil_align,
@@ -2342,7 +2349,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def prepare_for_extend(self):
self.forward_mode = ForwardMode.EXTEND
server_args = get_server_args()
if self.is_dllm():
# For DLLM, we use a separate forward mode
@@ -2471,7 +2477,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
req.already_computed = seq_len
req.is_retracted = False
if server_args.enable_mamba_extra_buffer():
if mamba_extra_buffer_enabled():
track_entry = self._mamba_radix_cache_v2_req_prepare_for_extend(req)
mamba_track_mask_cpu.append(track_entry.track_mask)
mamba_track_indices_cpu.append(track_entry.track_index)
@@ -2576,7 +2582,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.extend_logprob_start_lens = extend_logprob_start_lens
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
if server_args.enable_mamba_extra_buffer():
if mamba_extra_buffer_enabled():
self.mamba_track_indices = torch.tensor(
mamba_track_indices_cpu,
dtype=torch.int64,
@@ -2662,7 +2668,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# In lazy mode, skip the swap — the second ping-pong slot is not
# allocated yet; it will be allocated on demand at the track boundary
# in mamba_lazy_prealloc_at_boundary during prepare_for_decode.
if not server_args.enable_mamba_extra_buffer_lazy():
if not mamba_extra_buffer_lazy_enabled():
req.mamba_next_track_idx = (
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
req.mamba_next_track_idx
@@ -3003,7 +3009,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def prepare_for_decode(self):
self.forward_mode = ForwardMode.DECODE
server_args = get_server_args()
# Decode embeds the last output token via embed_tokens; clear the stale
# prefill-time tensor so it doesn't leak into ForwardBatch.
self.input_embeds = None
@@ -3057,7 +3062,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.req_pool_indices_cpu,
)
if server_args.enable_mamba_extra_buffer():
if mamba_extra_buffer_enabled():
mamba_track_interval = get_exec().mamba.mamba_track_interval
if len(self.reqs) == 0:
@@ -3065,7 +3070,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
(0,), dtype=torch.int64, device=self.device
)
else:
if server_args.enable_mamba_extra_buffer_lazy():
if mamba_extra_buffer_lazy_enabled():
self.mamba_lazy_prealloc_at_boundary(mamba_track_interval)
set_mamba_track_indices_from_reqs(self)
@@ -38,6 +38,7 @@ from sglang.srt.runtime_context import (
get_memory,
get_observability,
get_server_args,
mamba_extra_buffer_lazy_enabled,
)
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
@@ -1073,7 +1074,7 @@ class SchedulerBatchResultProcessor:
prepare_release(req)
is_insert = (
req.mamba_lazy_is_insert
if get_server_args().enable_mamba_extra_buffer_lazy()
if mamba_extra_buffer_lazy_enabled()
else True
)
release_kv_cache(req, self.tree_cache, is_insert=is_insert)
@@ -1109,7 +1110,7 @@ class SchedulerBatchResultProcessor:
if req.mamba_ping_pong_track_buffer is None:
return
lazy = get_server_args().enable_mamba_extra_buffer_lazy()
lazy = mamba_extra_buffer_lazy_enabled()
if known_boundary:
self._mamba_assert_committed_len_lookahead(req)
track_seqlen = req.kv_committed_len
@@ -6,10 +6,7 @@ from sglang.srt.runtime_context import get_server_args
from sglang.srt.server_args import ServerArgs
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if server_args is None:
server_args = get_server_args()
def get_alloc_len_per_decode(server_args: ServerArgs) -> int:
if server_args.speculative_algorithm is None:
return 1
@@ -39,7 +36,13 @@ def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> in
The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
mode; see eagle_utils.eagle_prepare_for_decode.
Callers on a request path have no config in hand, so this is the module's
single "which config" decision point: everything below it takes the
instance explicitly.
"""
if server_args is None:
server_args = get_server_args()
return 2 * get_alloc_len_per_decode(server_args)
+3 -8
View File
@@ -858,11 +858,7 @@ class DeepseekV2MoE(nn.Module):
)
]
def _can_dual_stream_graph(
self, hidden_states: torch.Tensor, server_args=None
) -> bool:
if server_args is None:
server_args = get_server_args()
def _can_dual_stream_graph(self, hidden_states: torch.Tensor) -> bool:
return (
_enable_pcg_dsv2_dual_stream
and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
@@ -875,7 +871,7 @@ class DeepseekV2MoE(nn.Module):
and not self._enable_a2a_moe
and not self._fuse_shared_experts_inside_sbo
and not getattr(self, "is_hash", False)
and not server_args.enable_eplb
and not get_exec().moe.enable_eplb
)
def forward(
@@ -898,8 +894,7 @@ class DeepseekV2MoE(nn.Module):
)
if not self._enable_a2a_moe:
server_args = get_server_args()
if self._can_dual_stream_graph(hidden_states, server_args):
if self._can_dual_stream_graph(hidden_states):
fwd = get_forward()
return dsv2_flashinfer_moe_dual_stream_graph(
hidden_states,
+2 -3
View File
@@ -80,7 +80,7 @@ from sglang.srt.runtime_context import (
get_model,
get_parallel,
get_schedule,
get_server_args,
mamba_extra_buffer_enabled,
)
from sglang.srt.utils import add_prefix, is_cuda, make_layers
@@ -1017,12 +1017,11 @@ class InklingForConditionalGeneration(nn.Module):
self.config = config
self.text_config = config.text_config
server_args = get_server_args()
assert envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get()
if get_disagg().disaggregation_mode != "decode":
assert not get_memory().disable_radix_cache
assert not get_schedule().disable_hybrid_swa_memory
assert server_args.enable_mamba_extra_buffer()
assert mamba_extra_buffer_enabled()
from types import SimpleNamespace
@@ -17,7 +17,11 @@ from sglang.srt.models.inkling_common.kernels.sconv import (
save_intermediate_conv_windows,
update_sconv_cache,
)
from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args
from sglang.srt.runtime_context import (
get_exec,
get_parallel,
mamba_extra_buffer_enabled,
)
from sglang.srt.utils import is_cuda, set_weight_attrs
@@ -267,10 +271,7 @@ class ShortConvolution(nn.Module):
draft_token_num = hidden_states.shape[1]
mamba_track_indices = getattr(forward_batch, "mamba_track_indices", None)
do_tracking = (
mamba_track_indices is not None
and get_server_args().enable_mamba_extra_buffer()
)
do_tracking = mamba_track_indices is not None and mamba_extra_buffer_enabled()
crossed = track_step = None
if do_tracking:
+23
View File
@@ -1383,3 +1383,26 @@ def reset_context() -> None:
_CONTEXT.resources = Resources()
_CONTEXT.forward = ForwardFlags()
set_global_dwdp_manager(None)
def mamba_extra_buffer_enabled() -> bool:
"""Whether the mamba radix cache keeps its extra state buffer.
A predicate over two published leaves (``memory.disable_radix_cache`` and
``exec.mamba.mamba_radix_cache_strategy``), so it reads the bags rather
than the startup record the ``ServerArgs`` member of the same name is the
pre-publish equivalent used inside the resolution pipeline.
"""
return (
get_memory().disable_radix_cache is False
and get_exec().mamba.mamba_radix_cache_strategy
in ("extra_buffer", "extra_buffer_lazy")
)
def mamba_extra_buffer_lazy_enabled() -> bool:
"""The lazy variant of :func:`mamba_extra_buffer_enabled`."""
return (
get_memory().disable_radix_cache is False
and get_exec().mamba.mamba_radix_cache_strategy == "extra_buffer_lazy"
)
+14 -11
View File
@@ -46,7 +46,13 @@ from sglang.srt.mem_cache.allocation import (
from sglang.srt.mem_cache.allocation import (
assign_req_to_token_pool_func as assign_req_to_token_pool_func,
)
from sglang.srt.runtime_context import get_exec, get_server_args
from sglang.srt.runtime_context import (
get_exec,
get_server_args,
get_spec,
mamba_extra_buffer_enabled,
mamba_extra_buffer_lazy_enabled,
)
from sglang.srt.utils import (
is_cpu,
is_cuda,
@@ -250,16 +256,14 @@ def record_stream_for_v2_verify(batch, verify_input, fwd_stream):
record_stream_each(candidates, fwd_stream)
def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
if server_args is None:
server_args = get_server_args()
def spec_need_hidden_states() -> bool:
# STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
# multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap.
# TODO(lsyin): also skip when step == 1.
if server_args.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"):
spec = get_spec()
if spec.speculative_algorithm in ("STANDALONE", "DFLASH", "DSPARK"):
return False
return not server_args.enable_multi_layer_eagle
return not spec.enable_multi_layer_eagle
@torch.compile(dynamic=True, disable=_is_npu or _is_xpu)
@@ -761,11 +765,10 @@ def prepare_mamba_track_for_verify(batch: ScheduleBatch) -> None:
Lazy: gather the positions planned by mamba_lazy_spec_prepare. Runs
inside forward isolation, so it must not mutate req/pool state.
"""
server_args = get_server_args()
if not server_args.enable_mamba_extra_buffer():
if not mamba_extra_buffer_enabled():
return
track_positions = None
if server_args.enable_mamba_extra_buffer_lazy():
if mamba_extra_buffer_lazy_enabled():
track_positions = batch.mamba_lazy_spec_track_positions_cpu
assert track_positions is not None and len(track_positions) == len(
batch.reqs
@@ -1028,7 +1031,7 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
prep on its draft input -- the dispatcher routes.
"""
server_args = get_server_args()
if server_args.enable_mamba_extra_buffer_lazy():
if mamba_extra_buffer_lazy_enabled():
# Scheduler phase (outside forward isolation).
batch.mamba_lazy_spec_prepare(
get_exec().mamba.mamba_track_interval,