config: retire the hidden global fallbacks and the mamba-extra-buffer instance reads (#34080)
This commit is contained in:
@@ -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]]]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user