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:
|
def mamba_extra_buffer_of(cfg: Any) -> bool:
|
||||||
"""Mid-resolution equivalent of runtime_context.mamba_extra_buffer_enabled:
|
"""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 (
|
return cfg.disable_radix_cache is False and cfg.mamba_radix_cache_strategy in (
|
||||||
"extra_buffer",
|
"extra_buffer",
|
||||||
"extra_buffer_lazy",
|
"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(
|
def collect_model_override_declarations(
|
||||||
architecture: str, server_args: Any, hf_config: Any
|
architecture: str, server_args: Any, hf_config: Any
|
||||||
) -> List[Tuple[str, Dict[str, 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.dllm.config import DllmConfig
|
||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
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 (
|
from sglang.srt.utils.common import (
|
||||||
Range,
|
Range,
|
||||||
ceil_align,
|
ceil_align,
|
||||||
@@ -2342,7 +2349,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
|
|
||||||
def prepare_for_extend(self):
|
def prepare_for_extend(self):
|
||||||
self.forward_mode = ForwardMode.EXTEND
|
self.forward_mode = ForwardMode.EXTEND
|
||||||
server_args = get_server_args()
|
|
||||||
|
|
||||||
if self.is_dllm():
|
if self.is_dllm():
|
||||||
# For DLLM, we use a separate forward mode
|
# For DLLM, we use a separate forward mode
|
||||||
@@ -2471,7 +2477,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
req.already_computed = seq_len
|
req.already_computed = seq_len
|
||||||
req.is_retracted = False
|
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)
|
track_entry = self._mamba_radix_cache_v2_req_prepare_for_extend(req)
|
||||||
mamba_track_mask_cpu.append(track_entry.track_mask)
|
mamba_track_mask_cpu.append(track_entry.track_mask)
|
||||||
mamba_track_indices_cpu.append(track_entry.track_index)
|
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_logprob_start_lens = extend_logprob_start_lens
|
||||||
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
|
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(
|
self.mamba_track_indices = torch.tensor(
|
||||||
mamba_track_indices_cpu,
|
mamba_track_indices_cpu,
|
||||||
dtype=torch.int64,
|
dtype=torch.int64,
|
||||||
@@ -2662,7 +2668,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
# In lazy mode, skip the swap — the second ping-pong slot is not
|
# 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
|
# allocated yet; it will be allocated on demand at the track boundary
|
||||||
# in mamba_lazy_prealloc_at_boundary during prepare_for_decode.
|
# 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 = (
|
req.mamba_next_track_idx = (
|
||||||
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
|
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
|
||||||
req.mamba_next_track_idx
|
req.mamba_next_track_idx
|
||||||
@@ -3003,7 +3009,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
|
|
||||||
def prepare_for_decode(self):
|
def prepare_for_decode(self):
|
||||||
self.forward_mode = ForwardMode.DECODE
|
self.forward_mode = ForwardMode.DECODE
|
||||||
server_args = get_server_args()
|
|
||||||
# Decode embeds the last output token via embed_tokens; clear the stale
|
# Decode embeds the last output token via embed_tokens; clear the stale
|
||||||
# prefill-time tensor so it doesn't leak into ForwardBatch.
|
# prefill-time tensor so it doesn't leak into ForwardBatch.
|
||||||
self.input_embeds = None
|
self.input_embeds = None
|
||||||
@@ -3057,7 +3062,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
self.req_pool_indices_cpu,
|
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
|
mamba_track_interval = get_exec().mamba.mamba_track_interval
|
||||||
|
|
||||||
if len(self.reqs) == 0:
|
if len(self.reqs) == 0:
|
||||||
@@ -3065,7 +3070,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
(0,), dtype=torch.int64, device=self.device
|
(0,), dtype=torch.int64, device=self.device
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
if server_args.enable_mamba_extra_buffer_lazy():
|
if mamba_extra_buffer_lazy_enabled():
|
||||||
self.mamba_lazy_prealloc_at_boundary(mamba_track_interval)
|
self.mamba_lazy_prealloc_at_boundary(mamba_track_interval)
|
||||||
set_mamba_track_indices_from_reqs(self)
|
set_mamba_track_indices_from_reqs(self)
|
||||||
|
|
||||||
|
|||||||
@@ -38,6 +38,7 @@ from sglang.srt.runtime_context import (
|
|||||||
get_memory,
|
get_memory,
|
||||||
get_observability,
|
get_observability,
|
||||||
get_server_args,
|
get_server_args,
|
||||||
|
mamba_extra_buffer_lazy_enabled,
|
||||||
)
|
)
|
||||||
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
|
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
|
||||||
from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
|
from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
|
||||||
@@ -1073,7 +1074,7 @@ class SchedulerBatchResultProcessor:
|
|||||||
prepare_release(req)
|
prepare_release(req)
|
||||||
is_insert = (
|
is_insert = (
|
||||||
req.mamba_lazy_is_insert
|
req.mamba_lazy_is_insert
|
||||||
if get_server_args().enable_mamba_extra_buffer_lazy()
|
if mamba_extra_buffer_lazy_enabled()
|
||||||
else True
|
else True
|
||||||
)
|
)
|
||||||
release_kv_cache(req, self.tree_cache, is_insert=is_insert)
|
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:
|
if req.mamba_ping_pong_track_buffer is None:
|
||||||
return
|
return
|
||||||
|
|
||||||
lazy = get_server_args().enable_mamba_extra_buffer_lazy()
|
lazy = mamba_extra_buffer_lazy_enabled()
|
||||||
if known_boundary:
|
if known_boundary:
|
||||||
self._mamba_assert_committed_len_lookahead(req)
|
self._mamba_assert_committed_len_lookahead(req)
|
||||||
track_seqlen = req.kv_committed_len
|
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
|
from sglang.srt.server_args import ServerArgs
|
||||||
|
|
||||||
|
|
||||||
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
|
def get_alloc_len_per_decode(server_args: ServerArgs) -> int:
|
||||||
if server_args is None:
|
|
||||||
server_args = get_server_args()
|
|
||||||
|
|
||||||
if server_args.speculative_algorithm is None:
|
if server_args.speculative_algorithm is None:
|
||||||
return 1
|
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
|
The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
|
||||||
mode; see eagle_utils.eagle_prepare_for_decode.
|
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)
|
return 2 * get_alloc_len_per_decode(server_args)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -858,11 +858,7 @@ class DeepseekV2MoE(nn.Module):
|
|||||||
)
|
)
|
||||||
]
|
]
|
||||||
|
|
||||||
def _can_dual_stream_graph(
|
def _can_dual_stream_graph(self, hidden_states: torch.Tensor) -> bool:
|
||||||
self, hidden_states: torch.Tensor, server_args=None
|
|
||||||
) -> bool:
|
|
||||||
if server_args is None:
|
|
||||||
server_args = get_server_args()
|
|
||||||
return (
|
return (
|
||||||
_enable_pcg_dsv2_dual_stream
|
_enable_pcg_dsv2_dual_stream
|
||||||
and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph())
|
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._enable_a2a_moe
|
||||||
and not self._fuse_shared_experts_inside_sbo
|
and not self._fuse_shared_experts_inside_sbo
|
||||||
and not getattr(self, "is_hash", False)
|
and not getattr(self, "is_hash", False)
|
||||||
and not server_args.enable_eplb
|
and not get_exec().moe.enable_eplb
|
||||||
)
|
)
|
||||||
|
|
||||||
def forward(
|
def forward(
|
||||||
@@ -898,8 +894,7 @@ class DeepseekV2MoE(nn.Module):
|
|||||||
)
|
)
|
||||||
|
|
||||||
if not self._enable_a2a_moe:
|
if not self._enable_a2a_moe:
|
||||||
server_args = get_server_args()
|
if self._can_dual_stream_graph(hidden_states):
|
||||||
if self._can_dual_stream_graph(hidden_states, server_args):
|
|
||||||
fwd = get_forward()
|
fwd = get_forward()
|
||||||
return dsv2_flashinfer_moe_dual_stream_graph(
|
return dsv2_flashinfer_moe_dual_stream_graph(
|
||||||
hidden_states,
|
hidden_states,
|
||||||
|
|||||||
@@ -80,7 +80,7 @@ from sglang.srt.runtime_context import (
|
|||||||
get_model,
|
get_model,
|
||||||
get_parallel,
|
get_parallel,
|
||||||
get_schedule,
|
get_schedule,
|
||||||
get_server_args,
|
mamba_extra_buffer_enabled,
|
||||||
)
|
)
|
||||||
from sglang.srt.utils import add_prefix, is_cuda, make_layers
|
from sglang.srt.utils import add_prefix, is_cuda, make_layers
|
||||||
|
|
||||||
@@ -1017,12 +1017,11 @@ class InklingForConditionalGeneration(nn.Module):
|
|||||||
self.config = config
|
self.config = config
|
||||||
self.text_config = config.text_config
|
self.text_config = config.text_config
|
||||||
|
|
||||||
server_args = get_server_args()
|
|
||||||
assert envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get()
|
assert envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get()
|
||||||
if get_disagg().disaggregation_mode != "decode":
|
if get_disagg().disaggregation_mode != "decode":
|
||||||
assert not get_memory().disable_radix_cache
|
assert not get_memory().disable_radix_cache
|
||||||
assert not get_schedule().disable_hybrid_swa_memory
|
assert not get_schedule().disable_hybrid_swa_memory
|
||||||
assert server_args.enable_mamba_extra_buffer()
|
assert mamba_extra_buffer_enabled()
|
||||||
|
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
|||||||
@@ -17,7 +17,11 @@ from sglang.srt.models.inkling_common.kernels.sconv import (
|
|||||||
save_intermediate_conv_windows,
|
save_intermediate_conv_windows,
|
||||||
update_sconv_cache,
|
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
|
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]
|
draft_token_num = hidden_states.shape[1]
|
||||||
|
|
||||||
mamba_track_indices = getattr(forward_batch, "mamba_track_indices", None)
|
mamba_track_indices = getattr(forward_batch, "mamba_track_indices", None)
|
||||||
do_tracking = (
|
do_tracking = mamba_track_indices is not None and mamba_extra_buffer_enabled()
|
||||||
mamba_track_indices is not None
|
|
||||||
and get_server_args().enable_mamba_extra_buffer()
|
|
||||||
)
|
|
||||||
|
|
||||||
crossed = track_step = None
|
crossed = track_step = None
|
||||||
if do_tracking:
|
if do_tracking:
|
||||||
|
|||||||
@@ -1383,3 +1383,26 @@ def reset_context() -> None:
|
|||||||
_CONTEXT.resources = Resources()
|
_CONTEXT.resources = Resources()
|
||||||
_CONTEXT.forward = ForwardFlags()
|
_CONTEXT.forward = ForwardFlags()
|
||||||
set_global_dwdp_manager(None)
|
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 (
|
from sglang.srt.mem_cache.allocation import (
|
||||||
assign_req_to_token_pool_func as assign_req_to_token_pool_func,
|
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 (
|
from sglang.srt.utils import (
|
||||||
is_cpu,
|
is_cpu,
|
||||||
is_cuda,
|
is_cuda,
|
||||||
@@ -250,16 +256,14 @@ def record_stream_for_v2_verify(batch, verify_input, fwd_stream):
|
|||||||
record_stream_each(candidates, fwd_stream)
|
record_stream_each(candidates, fwd_stream)
|
||||||
|
|
||||||
|
|
||||||
def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool:
|
def spec_need_hidden_states() -> bool:
|
||||||
if server_args is None:
|
|
||||||
server_args = get_server_args()
|
|
||||||
|
|
||||||
# STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
|
# STANDALONE drafts don't consume `spec_info.hidden_states` (vanilla LLM).
|
||||||
# multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap.
|
# multi_layer_eagle, DFLASH, and DSPARK don't relay hidden_states through FutureMap.
|
||||||
# TODO(lsyin): also skip when step == 1.
|
# 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 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)
|
@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
|
Lazy: gather the positions planned by mamba_lazy_spec_prepare. Runs
|
||||||
inside forward isolation, so it must not mutate req/pool state.
|
inside forward isolation, so it must not mutate req/pool state.
|
||||||
"""
|
"""
|
||||||
server_args = get_server_args()
|
if not mamba_extra_buffer_enabled():
|
||||||
if not server_args.enable_mamba_extra_buffer():
|
|
||||||
return
|
return
|
||||||
track_positions = None
|
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
|
track_positions = batch.mamba_lazy_spec_track_positions_cpu
|
||||||
assert track_positions is not None and len(track_positions) == len(
|
assert track_positions is not None and len(track_positions) == len(
|
||||||
batch.reqs
|
batch.reqs
|
||||||
@@ -1028,7 +1031,7 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
|
|||||||
prep on its draft input -- the dispatcher routes.
|
prep on its draft input -- the dispatcher routes.
|
||||||
"""
|
"""
|
||||||
server_args = get_server_args()
|
server_args = get_server_args()
|
||||||
if server_args.enable_mamba_extra_buffer_lazy():
|
if mamba_extra_buffer_lazy_enabled():
|
||||||
# Scheduler phase (outside forward isolation).
|
# Scheduler phase (outside forward isolation).
|
||||||
batch.mamba_lazy_spec_prepare(
|
batch.mamba_lazy_spec_prepare(
|
||||||
get_exec().mamba.mamba_track_interval,
|
get_exec().mamba.mamba_track_interval,
|
||||||
|
|||||||
@@ -1232,7 +1232,6 @@ class TestMlxOverlapScheduler(unittest.TestCase):
|
|||||||
logits_output = SimpleNamespace(customized_info=None)
|
logits_output = SimpleNamespace(customized_info=None)
|
||||||
original_release = batch_result_processor_module.release_kv_cache
|
original_release = batch_result_processor_module.release_kv_cache
|
||||||
original_get_indexer = batch_result_processor_module.get_global_indexer_capturer
|
original_get_indexer = batch_result_processor_module.get_global_indexer_capturer
|
||||||
original_get_server_args = batch_result_processor_module.get_server_args
|
|
||||||
|
|
||||||
def fake_release_kv_cache(release_req, tree_cache, is_insert=False):
|
def fake_release_kv_cache(release_req, tree_cache, is_insert=False):
|
||||||
events.append(("release", release_req.rid))
|
events.append(("release", release_req.rid))
|
||||||
@@ -1240,21 +1239,26 @@ class TestMlxOverlapScheduler(unittest.TestCase):
|
|||||||
|
|
||||||
batch_result_processor_module.release_kv_cache = fake_release_kv_cache
|
batch_result_processor_module.release_kv_cache = fake_release_kv_cache
|
||||||
batch_result_processor_module.get_global_indexer_capturer = lambda: None
|
batch_result_processor_module.get_global_indexer_capturer = lambda: None
|
||||||
batch_result_processor_module.get_server_args = lambda: SimpleNamespace(
|
# The lazy predicate reads the published bags; publish the non-lazy
|
||||||
enable_mamba_extra_buffer_lazy=lambda: False
|
# strategy instead of stubbing the accessor.
|
||||||
|
from sglang.srt.runtime_context import get_context
|
||||||
|
|
||||||
|
override = get_context().override_server_args(
|
||||||
|
mamba_radix_cache_strategy="extra_buffer"
|
||||||
)
|
)
|
||||||
|
override.install()
|
||||||
try:
|
try:
|
||||||
SchedulerBatchResultProcessor._handle_finish_state_updated_req(
|
SchedulerBatchResultProcessor._handle_finish_state_updated_req(
|
||||||
processor, req, batch, result, i, logits_output
|
processor, req, batch, result, i, logits_output
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
|
override.restore()
|
||||||
for name, original in saved.items():
|
for name, original in saved.items():
|
||||||
setattr(SchedulerBatchResultProcessor, name, original)
|
setattr(SchedulerBatchResultProcessor, name, original)
|
||||||
batch_result_processor_module.release_kv_cache = original_release
|
batch_result_processor_module.release_kv_cache = original_release
|
||||||
batch_result_processor_module.get_global_indexer_capturer = (
|
batch_result_processor_module.get_global_indexer_capturer = (
|
||||||
original_get_indexer
|
original_get_indexer
|
||||||
)
|
)
|
||||||
batch_result_processor_module.get_server_args = original_get_server_args
|
|
||||||
|
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
events,
|
events,
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from sglang.srt.managers.scheduler import Scheduler
|
|||||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||||
SchedulerBatchResultProcessor,
|
SchedulerBatchResultProcessor,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.runtime_context import get_context
|
||||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
|
||||||
@@ -144,42 +145,25 @@ class TestMambaBoundaryMaskReuse(unittest.TestCase):
|
|||||||
processor.process_batch_result_decode(result_batch, batch_result)
|
processor.process_batch_result_decode(result_batch, batch_result)
|
||||||
|
|
||||||
scheduler.process_batch_result = process_batch_result
|
scheduler.process_batch_result = process_batch_result
|
||||||
server_args = SimpleNamespace(
|
|
||||||
enable_mamba_extra_buffer=lambda: True,
|
|
||||||
enable_mamba_extra_buffer_lazy=lambda: False,
|
|
||||||
)
|
|
||||||
|
|
||||||
with (
|
with (
|
||||||
|
# The mamba predicates and the track interval read the
|
||||||
|
# published bags, so publish the configuration under test
|
||||||
|
# (non-lazy extra buffer, interval 4); observability and
|
||||||
|
# disagg reads are served by the same publish at their
|
||||||
|
# defaults.
|
||||||
|
get_context().override_server_args(
|
||||||
|
mamba_radix_cache_strategy="extra_buffer",
|
||||||
|
mamba_track_interval=4,
|
||||||
|
),
|
||||||
patch(
|
patch(
|
||||||
"sglang.srt.managers.schedule_batch.alloc_for_decode",
|
"sglang.srt.managers.schedule_batch.alloc_for_decode",
|
||||||
return_value=torch.tensor([3], dtype=torch.int64),
|
return_value=torch.tensor([3], dtype=torch.int64),
|
||||||
),
|
),
|
||||||
patch(
|
|
||||||
"sglang.srt.managers.schedule_batch.get_server_args",
|
|
||||||
return_value=server_args,
|
|
||||||
),
|
|
||||||
patch(
|
|
||||||
"sglang.srt.managers.schedule_batch.get_exec",
|
|
||||||
return_value=SimpleNamespace(
|
|
||||||
mamba=SimpleNamespace(mamba_track_interval=4)
|
|
||||||
),
|
|
||||||
),
|
|
||||||
patch(
|
patch(
|
||||||
"sglang.srt.managers.schedule_batch.set_mamba_track_indices_from_reqs"
|
"sglang.srt.managers.schedule_batch.set_mamba_track_indices_from_reqs"
|
||||||
),
|
),
|
||||||
patch.object(torch.Tensor, "pin_memory", lambda tensor: tensor),
|
patch.object(torch.Tensor, "pin_memory", lambda tensor: tensor),
|
||||||
patch(
|
|
||||||
"sglang.srt.managers.scheduler_components."
|
|
||||||
"batch_result_processor.get_observability",
|
|
||||||
return_value=SimpleNamespace(enable_metrics=False),
|
|
||||||
),
|
|
||||||
patch(
|
|
||||||
"sglang.srt.managers.scheduler_components."
|
|
||||||
"batch_result_processor.get_disagg",
|
|
||||||
return_value=SimpleNamespace(
|
|
||||||
disaggregation_decode_enable_offload_kvcache=False
|
|
||||||
),
|
|
||||||
),
|
|
||||||
patch.object(
|
patch.object(
|
||||||
SchedulerBatchResultProcessor,
|
SchedulerBatchResultProcessor,
|
||||||
"_mamba_prefix_cache_update",
|
"_mamba_prefix_cache_update",
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from unittest.mock import patch
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.runtime_context import get_context
|
||||||
from sglang.test.ci.ci_register import register_cpu_ci
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
from sglang.test.test_utils import maybe_stub_sgl_kernel
|
from sglang.test.test_utils import maybe_stub_sgl_kernel
|
||||||
|
|
||||||
@@ -43,18 +44,16 @@ class TestPrepareForDecodeSeqLensOwnership(unittest.TestCase):
|
|||||||
"""Each prepare_for_decode call rebinds seq-lens tensors to new +1 objects without mutating the old ones."""
|
"""Each prepare_for_decode call rebinds seq-lens tensors to new +1 objects without mutating the old ones."""
|
||||||
batch = _make_decode_batch()
|
batch = _make_decode_batch()
|
||||||
|
|
||||||
server_args = types.SimpleNamespace(
|
# The mamba-extra-buffer predicate reads the published bags, so the
|
||||||
enable_mamba_extra_buffer=lambda: False,
|
# fixture publishes a config with the strategy off.
|
||||||
|
override = get_context().override_server_args(
|
||||||
|
mamba_radix_cache_strategy="no_buffer"
|
||||||
)
|
)
|
||||||
with (
|
override.install()
|
||||||
patch(
|
self.addCleanup(override.restore)
|
||||||
|
with patch(
|
||||||
"sglang.srt.managers.schedule_batch.alloc_for_decode",
|
"sglang.srt.managers.schedule_batch.alloc_for_decode",
|
||||||
return_value=torch.tensor([6, 7], dtype=torch.int64),
|
return_value=torch.tensor([6, 7], dtype=torch.int64),
|
||||||
),
|
|
||||||
patch(
|
|
||||||
"sglang.srt.managers.schedule_batch.get_server_args",
|
|
||||||
return_value=server_args,
|
|
||||||
),
|
|
||||||
):
|
):
|
||||||
for step in range(1, 3):
|
for step in range(1, 3):
|
||||||
prev_seq_lens = batch.seq_lens
|
prev_seq_lens = batch.seq_lens
|
||||||
|
|||||||
Reference in New Issue
Block a user