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

[Unified Tree] Support Branching-Point Caching for the SWA Component (#34565)

Browse Source 
Co-authored-by: alphabetc1 <2508695655@qq.com>
Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>
This commit is contained in:
Jincong Chen
2026-09-05 12:29:16 +08:00
committed by GitHub
co-authored by alphabetc1 Shuwen Wang
parent 991368d880
commit 3a770da756
10 changed files with 467 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+3
View File
@@ -117,6 +117,9 @@ class ScheduleBatchDisaggregationDecodeMixin:
last_tokens: List[int] = []
for req in self.reqs:
last_tokens.append(req.output_ids[-1])
# PREBUILT does not materialize a local SWA branching window.
if req.swa_branching_seqlen is not None:
req.swa_branching_seqlen = None
maybe_cache_unfinished_req(req, self.tree_cache)
if req.grammar is not None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
python/sglang/srt/managers/schedule_batch.py
+7
View File
@@ -1004,6 +1004,10 @@ class Req(ReqDllmMixin):
# For req-level memory management
self.kv = ReqKvInfo()
# Full-KV-derived boundary whose SWA window should be inserted after
# the current prefill pass.
self.swa_branching_seqlen: Optional[int] = None
# for cross-encoder model
self.token_type_ids = token_type_ids
@@ -1515,6 +1519,7 @@ class Req(ReqDllmMixin):
self.best_match_node,
self.host_hit_length,
self.swa_host_hit_length,
self.swa_branching_seqlen,
self.mamba_host_hit_length,
self.mamba_branching_seqlen,
) = (
@@ -1524,6 +1529,7 @@ class Req(ReqDllmMixin):
match_result.best_match_node,
match_result.host_hit_length,
match_result.swa_host_hit_length,
match_result.swa_branching_seqlen,
match_result.mamba_host_hit_length,
match_result.mamba_branching_seqlen,
)
@@ -1817,6 +1823,7 @@ class Req(ReqDllmMixin):
self.num_matched_prefix_tokens = 0
self.swa_uuid_for_lock = None
self.swa_prefix_lock_released = False
self.swa_branching_seqlen = None
self.skip_lock_node_ids = {}
self.extend_range = None
self.dllm_initialized = False
python/sglang/srt/managers/schedule_policy.py
+1
View File
@@ -195,6 +195,7 @@ def match_prefix_for_req(
req.num_matched_prefix_tokens = min(
len(req.prefix_indices) + req.host_hit_length, max_len
)
req.swa_branching_seqlen = match_result.swa_branching_seqlen
if match_result.mamba_branching_seqlen is not None:
req.mamba_branching_seqlen = match_result.mamba_branching_seqlen
if match_result.cache_protected_len is not None:
python/sglang/srt/mem_cache/base_prefix_cache.py
+7
View File
@@ -73,6 +73,7 @@ class InsertParams:
# SWA specific
prev_prefix_len: int = 0
swa_evicted_seqlen: int = 0
swa_branching_seqlen: Optional[int] = None
# General
chunked: bool = False
@@ -88,6 +89,7 @@ class InsertResult:
total_len: int = 0
last_device_node: Any = None
mamba_exist: bool = False
swa_branch_inserted: bool = False
inserted_host_node: Any = None
host_insert_dropped: bool = False
adopted_ranges: Optional[dict[ComponentType, list[tuple[int, int]]]] = None
@@ -201,6 +203,9 @@ class MatchResult(NamedTuple):
loaded back to device. Pure-KV cache semantics;
swa_host_hit_length : Number of SWA tokens that hit on host (within the sliding
window) and will be load-back into the SWA device pool.
swa_branching_seqlen: The SWA radix cache branching point, which is the longest
page-aligned position that could've been cache hit if there
exists an SWA window.
mamba_host_hit_length: Number of Mamba slots that hit on host and will be load-back
into the Mamba device pool. Typically 0 or 1.
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
@@ -216,6 +221,7 @@ class MatchResult(NamedTuple):
best_match_node: Any
host_hit_length: int = 0
swa_host_hit_length: int = 0
swa_branching_seqlen: Optional[int] = None
mamba_host_hit_length: int = 0
mamba_branching_seqlen: Optional[int] = None
cache_protected_len: Optional[int] = None
@@ -240,6 +246,7 @@ def zero_match_result(
best_match_node=root,
host_hit_length=0,
swa_host_hit_length=0,
swa_branching_seqlen=None,
mamba_host_hit_length=0,
full_kv_hit_length=0,
)
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+117 -26
View File
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Sequence
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
@@ -84,8 +85,32 @@ class SWAComponent(TreeComponent):
component_type = ComponentType.SWA
def _dirty_backup_window(self, node: UnifiedTreeNode) -> list[UnifiedTreeNode]:
if not self.tree_core.has_swa_host_pool:
return []
ct = self.component_type
covered = 0
dirty: list[UnifiedTreeNode] = []
cur = node
while (
cur is not self.tree_core.root_node and covered < self.sliding_window_size
):
if cur.write_through_pending_id is not None:
break
cd = cur.component_data[ct]
value = cd.value if cd.value is not None else cd.host_value
if value is None:
break
covered += len(value)
if cd.value is not None and cd.host_value is None:
dirty.append(cur)
cur = cur.parent
return dirty
def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool:
return False
return bool(self._dirty_backup_window(node))
def reset_session_state(self) -> None:
super().reset_session_state()
@@ -308,6 +333,16 @@ class SWAComponent(TreeComponent):
best_value_len: int,
) -> MatchResult:
ct = self.component_type
swa_boundary_len = len(result.device_indices) + result.host_hit_length
# Full KV may extend beyond the latest reusable SWA window. The branching
# point is the last page-aligned position within the Full-KV hit that lies
# beyond the current SWA boundary.
aligned_seqlen = (
result.full_kv_hit_length // self.tree_core.page_size
) * self.tree_core.page_size
branching_seqlen = aligned_seqlen if aligned_seqlen > swa_boundary_len else None
n_swa = 0
swa_host_hit = 0
node = result.best_match_node
@@ -328,11 +363,11 @@ class SWAComponent(TreeComponent):
else:
break
node = node.parent
if swa_host_hit > 0:
return result._replace(
swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit)
)
return result
return result._replace(
swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit),
swa_branching_seqlen=branching_seqlen,
)
def update_component_on_insert_overlap(
self,
@@ -466,6 +501,11 @@ class SWAComponent(TreeComponent):
result: InsertResult,
cache_actions: list[CacheAction | ComponentAction],
) -> None:
branching_seqlen = params.swa_branching_seqlen
if branching_seqlen is not None:
assert params.key is not None
result.swa_branch_inserted = len(params.key) >= branching_seqlen
if not is_new_leaf:
return
@@ -857,23 +897,59 @@ class SWAComponent(TreeComponent):
# Unfinished requests can already have an SWA-evicted prefix; preserve
# that boundary so insertion creates a tombstone instead of live SWA KV.
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
return None
branching_seqlen = req.swa_branching_seqlen
if branching_seqlen is None or branching_seqlen <= req.kv.cache_protected_len:
return None
# An EAGLE key with N bigrams spans N + 1 raw tokens.
effective_cache_len = branching_seqlen + int(self.tree_core.is_eagle)
if effective_cache_len > token_ids_len:
return None
# Record the logical SWA branch boundary for insertion.
insert_params.swa_branching_seqlen = branching_seqlen
return effective_cache_len
def _free_out_of_window_slots(self, req: Req, pre_len: int) -> None:
if self.sliding_window_size is None:
return
free_swa_out_of_window_slots(
req,
pre_len,
sliding_window_size=self.sliding_window_size,
page_size=self.cache.page_size,
req_to_token_pool=self.cache.req_to_token_pool,
token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
retain_floor=self.cache.swa_retain_floor(req),
)
def free_out_of_window_slots(
self, req: Req, pre_len: int, insert_params: InsertParams
) -> None:
if self.sliding_window_size is not None:
free_swa_out_of_window_slots(
req,
pre_len,
sliding_window_size=self.sliding_window_size,
page_size=self.cache.page_size,
req_to_token_pool=self.cache.req_to_token_pool,
token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
retain_floor=self.cache.swa_retain_floor(req),
)
self._free_out_of_window_slots(req, pre_len)
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
def cleanup_after_caching_req(
self,
req: Req,
is_finished: bool,
insert_result: Optional[InsertResult] = None,
insert_params: Optional[InsertParams] = None,
) -> None:
if insert_result is not None and insert_result.swa_branch_inserted:
req.swa_branching_seqlen = None
# Free unused SWA slots after inserting the branch.
if (
not is_finished
and insert_result is not None
and insert_result.swa_branch_inserted
and envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get()
):
forward_key_len = len(req.get_fill_ids()) - int(self.tree_core.is_eagle)
self._free_out_of_window_slots(req, forward_key_len - 1)
# ---- HiCache Hooks ----
def prepare_prefetch(
@@ -933,15 +1009,22 @@ class SWAComponent(TreeComponent):
return None
if phase == CacheTransferPhase.BACKUP_HOST:
cd = node.component_data[ct]
if cd.value is None:
if self.cache.host_memory_mode == "buffer_only":
# Buffer mode stages one node/hash span per FIFO backup intent.
cd = node.component_data[ct]
dirty = [node] if cd.value is not None else []
else:
dirty = self._dirty_backup_window(node)
if not dirty:
return None
# cd.value already holds SWA-pool indices (translated at insert time).
# Host pool indexing wants int64.
dirty.reverse()
return [
PoolTransfer(
name=PoolName.SWA,
device_indices=cd.value.to(torch.int64),
device_indices=torch.cat(
[n.component_data[ct].value for n in dirty]
).to(torch.int64),
nodes_to_load=[n.id for n in dirty],
)
]
@@ -1120,10 +1203,18 @@ class SWAComponent(TreeComponent):
ct = self.component_type
if phase == CacheTransferPhase.BACKUP_HOST:
if transfers and transfers[0].host_indices is not None:
cd = node.component_data[ct]
if cd.host_value is None:
cd.host_value = transfers[0].host_indices.clone()
if not transfers or transfers[0].host_indices is None:
return
xfer = transfers[0]
offset = 0
for node_id in xfer.nodes_to_load or [node.id]:
target = self.tree_core.node_by_id(node_id)
cd = target.component_data[ct]
assert cd.value is not None and cd.host_value is None
size = len(cd.value)
cd.host_value = xfer.host_indices[offset : offset + size].clone()
offset += size
assert offset == len(xfer.host_indices)
return
if phase == CacheTransferPhase.LOAD_BACK:
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+7 -3
View File
@@ -1131,9 +1131,13 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
state.phase = _InsertPhase.TAIL
def _needs_incremental_component_backup(self, node: UnifiedTreeNode) -> bool:
components = self.components
if self.is_write_back:
swa = self.components_by_type.get(ComponentType.SWA)
components = () if swa is None else (swa,)
return any(
component.needs_incremental_backup(node)
for component in self.components
for component in components
if component.component_type != BASE_COMPONENT_TYPE
)
@@ -1147,7 +1151,6 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
node = state.target_node
return (
self.enable_hicache
and not self.is_write_back
and node.backuped
and node.write_through_pending_id is None
and self._needs_incremental_component_backup(node)
@@ -1979,7 +1982,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
for comp in self.components:
if comp.component_type == BASE_COMPONENT_TYPE:
continue
if node.component_data[comp.component_type].host_value is not None:
cd = node.component_data[comp.component_type]
if cd.host_value is not None and not comp.needs_incremental_backup(node):
continue
t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST)
if t:
python/sglang/srt/session/streaming_session.py
+2
View File
@@ -66,6 +66,8 @@ class SessionSlot:
# Later turns run on the slot's record (see restore_to_req).
assert kv is self.kv
req.swa_branching_seqlen = None
def restore_to_req(self, req: Req):
"""Restore KV state from this slot into an incoming request."""
req.kv = self.kv