minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Unified Cache][5/N]: Integrate external linker mode end to end (#37381)

Browse Source 
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
This commit is contained in:
Zhangheng
2026-09-04 02:02:58 +08:00
committed by GitHub
co-authored by 晟海
parent 619ab2bcce
commit abed680320
13 changed files with 586 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/hicache_hook.py
+13
View File
@@ -23,6 +23,19 @@ def handle_hicache(server_args: Any):
2) Storage <-> layout compatibility (may rewrite layout).
"""
cfg = resolving_view(server_args)
if cfg.enable_unified_cache_external_linker:
if cfg.enable_hierarchical_cache:
raise ValueError(
"--enable-unified-cache-external-linker and "
"--enable-hierarchical-cache are mutually exclusive."
)
if cfg.hicache_storage_backend is not None:
raise ValueError(
"--enable-unified-cache-external-linker does not use "
"--hicache-storage-backend."
)
return
# Skip all normalization when neither hicache nor decode-offload path is active.
if not (
cfg.enable_hierarchical_cache
python/sglang/srt/managers/scheduler.py
+28 -18
View File
@@ -481,6 +481,9 @@ class Scheduler(
self.enable_hierarchical_cache = get_memory().enable_hierarchical_cache
self.enable_session_radix_cache = get_memory().enable_session_radix_cache
self.enable_hicache_storage = get_memory().hicache_storage_backend is not None
self.enable_unified_cache_external_linker = (
get_memory().enable_unified_cache_external_linker
)
self.enable_decode_hicache = (
get_disagg().disaggregation_decode_enable_radix_cache
and self.enable_hierarchical_cache
@@ -3132,6 +3135,15 @@ class Scheduler(
return False
return True
def _release_aborted_request(self, rid: str) -> None:
"""Drop the cache-side state an aborted request left behind."""
if (
self.enable_hierarchical_cache
or self.enable_hicache_storage
or self.enable_unified_cache_external_linker
):
self.tree_cache.release_aborted_request(rid)
def _abort_on_queued_limit(self, recv_req: Req) -> bool:
"""Abort an incoming or existing request if the waiting queue is full. Returns True if the incoming request is aborted."""
if (
@@ -3158,9 +3170,7 @@ class Scheduler(
direction * recv_req.priority < direction * candidate_req.priority
)
if abort_existing_req:
if self.enable_hicache_storage:
# Release prefetch events associated with the request
self.tree_cache.release_aborted_request(candidate_req.rid)
self._release_aborted_request(candidate_req.rid)
self.waiting_queue.pop(idx)
self.beam_coordinator.retire_group(candidate_req)
req_to_abort = candidate_req
@@ -3189,9 +3199,7 @@ class Scheduler(
for req in self.waiting_queue:
entry_time = req.time_stats.wait_queue_entry_time
if 0 < entry_time < deadline:
if self.enable_hicache_storage:
# Release prefetch events associated with the request
self.tree_cache.release_aborted_request(req.rid)
self._release_aborted_request(req.rid)
self.ipc_channels.send_to_tokenizer.send_output(
_make_abort_req(
req,
@@ -3347,8 +3355,7 @@ class Scheduler(
req, self.req_to_metadata_buffer_idx_allocator
)
req.pending_bootstrap = False
if self.enable_hicache_storage:
self.tree_cache.release_aborted_request(req.rid)
self._release_aborted_request(req.rid)
release_kv_cache(req, self.tree_cache, is_insert=False)
self.chunked_req = None
@@ -3607,7 +3614,11 @@ class Scheduler(
for req in ready_grammar_requests:
self._add_request_to_queue(req)
if self.enable_hierarchical_cache or get_memory().enable_flexkv:
if (
self.enable_hierarchical_cache
or get_memory().enable_flexkv
or self.enable_unified_cache_external_linker
):
self.tree_cache.check_hicache_events()
if self.enable_hicache_storage:
self._retry_missed_storage_prefetches()
@@ -3780,7 +3791,10 @@ class Scheduler(
if res != AddReqResult.CONTINUE:
if res == AddReqResult.NO_TOKEN:
if self.enable_hierarchical_cache:
if (
self.enable_hierarchical_cache
or self.enable_unified_cache_external_linker
):
# Set batch_is_full after making sure there are requests that can be served
running_batch.batch_is_full = len(adder.can_run_list) > 0 or (
not running_batch.is_empty()
@@ -3844,7 +3858,7 @@ class Scheduler(
self.chunked_req is None or len(can_run_list) != 1
)
if self.enable_hierarchical_cache:
if self.enable_hierarchical_cache or self.enable_unified_cache_external_linker:
# todo (zhiqiang): disable cuda graph execution if hicache loading triggered
new_batch.hicache_consumer_index = (
self.tree_cache.ready_to_load_host_cache()
@@ -5052,10 +5066,8 @@ class Scheduler(
# This only works for requests that have not started anything.
# We still need to send something back to TokenizerManager to clean up the state.
req = self.waiting_queue.pop(i)
self._release_aborted_request(req.rid)
self.beam_coordinator.retire_group(req)
if self.enable_hicache_storage:
# to release prefetch events associated with the request
self.tree_cache.release_aborted_request(req.rid)
self.ipc_channels.send_to_tokenizer.send_output(_make_abort_req(req), req)
# For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
if self.disaggregation_mode == DisaggregationMode.DECODE:
@@ -5086,8 +5098,7 @@ class Scheduler(
for req in self.dllm_manager.pop_aborted_reqs(
recv_req.abort_all, recv_req.rid
):
if self.enable_hicache_storage:
self.tree_cache.release_aborted_request(req.rid)
self._release_aborted_request(req.rid)
self.ipc_channels.send_to_tokenizer.send_output(
_make_abort_req(req), req
)
@@ -5107,8 +5118,7 @@ class Scheduler(
for req in self.disagg_prefill_bootstrap_queue.queue:
if recv_req.abort_all or req.rid.startswith(recv_req.rid):
logger.debug(f"Abort bootstrap queue request. {req.rid=}")
if self.enable_hicache_storage:
self.tree_cache.release_aborted_request(req.rid)
self._release_aborted_request(req.rid)
if hasattr(req.disagg_kv_sender, "abort"):
req.disagg_kv_sender.abort()
python/sglang/srt/mem_cache/hicache_storage.py
+1 -1
View File
@@ -129,7 +129,7 @@ class PoolTransferResult:
extra_pool_hit_pages: dict[str, int]
# Pools with TRAILING_PAGES (SWA, Mamba state) only hold a window that ends on an
# offloaded node boundary.
# offloaded node boundary, so 5 can be restorable while 4 and 3 are not.
# Each rank owns its own shard and may hold a different set, so reducing a
# per-rank maximum would pick a length that is illegal on another rank; the
# caller intersects these sets instead.
python/sglang/srt/mem_cache/kv_cache_builder.py
+2
View File
@@ -313,6 +313,8 @@ def build_kv_cache(
enable_mamba_extra_buffer_lazy=server_args.enable_mamba_extra_buffer_lazy(),
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
attn_cp_rank=ps.attn_cp_rank,
attn_cp_size=ps.attn_cp_size,
chunked_prefill_size=effective_chunked_prefill_size,
sliding_window_size=sliding_window_size,
mtp_draft_device_pools=mtp_draft_device_pools,
python/sglang/srt/mem_cache/kv_cache_configurator.py
+1
View File
@@ -105,6 +105,7 @@ def _should_elide_dsa_index_k(*, is_draft_worker: bool) -> bool:
not memory_config.enable_hisparse
and not is_draft_worker
and not memory_config.enable_hierarchical_cache
and not memory_config.enable_unified_cache_external_linker
and get_disagg().disaggregation_mode == "null"
)
python/sglang/srt/mem_cache/registry.py
+23
View File
@@ -108,6 +108,9 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
logger.info("Using experimental C++ radix tree implementation.")
return RadixCacheCpp(params=params, server_args=server_args)
if server_args.enable_unified_cache_external_linker:
return _create_unified_radix_cache(ctx, server_args, params)
if ctx.is_hybrid_swa and ctx.full_tokens_per_layer == 0:
from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache
@@ -193,6 +196,26 @@ def _create_unified_radix_cache(
ctx.tp_worker.register_hicache_layer_transfer_counter(
cache.cache_controller.layer_done_counter
)
elif server_args.enable_unified_cache_external_linker:
backend = server_args.unified_cache_external_linker_backend
if backend == "mooncake":
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_direct_linker import (
MooncakeDirectLinker,
)
linker_cls = MooncakeDirectLinker
else:
raise ValueError(
f"Unknown unified cache external linker backend: {backend!r}"
)
cache.init_cache_linker(
linker_cls(server_args, params, components=set(cache.components))
)
counter = cache.linker.layer_done_counter
kvcache = params.token_to_kv_pool_allocator.get_kvcache()
kvcache.register_layer_transfer_counter(counter)
ctx.tp_worker.register_hicache_layer_transfer_counter(counter)
return cache
python/sglang/srt/mem_cache/unified_cache/components/full_component.py
+3
View File
@@ -487,7 +487,10 @@ class FullComponent(TreeComponent):
if phase == ExternalLinkerLoadPhase.ABORT:
self._full_allocator().free(transfer.device_indices)
return None
if phase == ExternalLinkerLoadPhase.PREPARE:
return transfer
assert phase == ExternalLinkerLoadPhase.COMMIT
return transfer
def free_host_values(self, host_values: list[torch.Tensor]) -> None:
python/sglang/srt/server_args.py
+17
View File
@@ -2821,6 +2821,23 @@ class ServerArgs:
NS("memory"),
] = 4
# -------------------------------------------------------------------------
# Unified Radix Cache
# -------------------------------------------------------------------------
enable_unified_cache_external_linker: A[
bool,
"Link UnifiedRadixCache directly to an external KV store (direct L3), with no host cache tier.",
NS("memory"),
] = False
unified_cache_external_linker_backend: A[
str,
Arg(
help="Storage backend for --enable-unified-cache-external-linker.",
choices=["mooncake"],
),
NS("memory"),
] = "mooncake"
# -------------------------------------------------------------------------
# Hierarchical sparse attention
# -------------------------------------------------------------------------