[Unified Tree] Support Branching-Point Caching for the SWA Component (#34565)
Co-authored-by: alphabetc1 <2508695655@qq.com> Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>
This commit is contained in:
co-authored by
alphabetc1
Shuwen Wang
parent
991368d880
commit
3a770da756
@@ -117,6 +117,9 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
|||||||
last_tokens: List[int] = []
|
last_tokens: List[int] = []
|
||||||
for req in self.reqs:
|
for req in self.reqs:
|
||||||
last_tokens.append(req.output_ids[-1])
|
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)
|
maybe_cache_unfinished_req(req, self.tree_cache)
|
||||||
if req.grammar is not None:
|
if req.grammar is not None:
|
||||||
# FIXME: this try-except block is for handling unexpected xgrammar issue.
|
# FIXME: this try-except block is for handling unexpected xgrammar issue.
|
||||||
|
|||||||
@@ -1004,6 +1004,10 @@ class Req(ReqDllmMixin):
|
|||||||
# For req-level memory management
|
# For req-level memory management
|
||||||
self.kv = ReqKvInfo()
|
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
|
# for cross-encoder model
|
||||||
self.token_type_ids = token_type_ids
|
self.token_type_ids = token_type_ids
|
||||||
|
|
||||||
@@ -1515,6 +1519,7 @@ class Req(ReqDllmMixin):
|
|||||||
self.best_match_node,
|
self.best_match_node,
|
||||||
self.host_hit_length,
|
self.host_hit_length,
|
||||||
self.swa_host_hit_length,
|
self.swa_host_hit_length,
|
||||||
|
self.swa_branching_seqlen,
|
||||||
self.mamba_host_hit_length,
|
self.mamba_host_hit_length,
|
||||||
self.mamba_branching_seqlen,
|
self.mamba_branching_seqlen,
|
||||||
) = (
|
) = (
|
||||||
@@ -1524,6 +1529,7 @@ class Req(ReqDllmMixin):
|
|||||||
match_result.best_match_node,
|
match_result.best_match_node,
|
||||||
match_result.host_hit_length,
|
match_result.host_hit_length,
|
||||||
match_result.swa_host_hit_length,
|
match_result.swa_host_hit_length,
|
||||||
|
match_result.swa_branching_seqlen,
|
||||||
match_result.mamba_host_hit_length,
|
match_result.mamba_host_hit_length,
|
||||||
match_result.mamba_branching_seqlen,
|
match_result.mamba_branching_seqlen,
|
||||||
)
|
)
|
||||||
@@ -1817,6 +1823,7 @@ class Req(ReqDllmMixin):
|
|||||||
self.num_matched_prefix_tokens = 0
|
self.num_matched_prefix_tokens = 0
|
||||||
self.swa_uuid_for_lock = None
|
self.swa_uuid_for_lock = None
|
||||||
self.swa_prefix_lock_released = False
|
self.swa_prefix_lock_released = False
|
||||||
|
self.swa_branching_seqlen = None
|
||||||
self.skip_lock_node_ids = {}
|
self.skip_lock_node_ids = {}
|
||||||
self.extend_range = None
|
self.extend_range = None
|
||||||
self.dllm_initialized = False
|
self.dllm_initialized = False
|
||||||
|
|||||||
@@ -195,6 +195,7 @@ def match_prefix_for_req(
|
|||||||
req.num_matched_prefix_tokens = min(
|
req.num_matched_prefix_tokens = min(
|
||||||
len(req.prefix_indices) + req.host_hit_length, max_len
|
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:
|
if match_result.mamba_branching_seqlen is not None:
|
||||||
req.mamba_branching_seqlen = match_result.mamba_branching_seqlen
|
req.mamba_branching_seqlen = match_result.mamba_branching_seqlen
|
||||||
if match_result.cache_protected_len is not None:
|
if match_result.cache_protected_len is not None:
|
||||||
|
|||||||
@@ -73,6 +73,7 @@ class InsertParams:
|
|||||||
# SWA specific
|
# SWA specific
|
||||||
prev_prefix_len: int = 0
|
prev_prefix_len: int = 0
|
||||||
swa_evicted_seqlen: int = 0
|
swa_evicted_seqlen: int = 0
|
||||||
|
swa_branching_seqlen: Optional[int] = None
|
||||||
|
|
||||||
# General
|
# General
|
||||||
chunked: bool = False
|
chunked: bool = False
|
||||||
@@ -88,6 +89,7 @@ class InsertResult:
|
|||||||
total_len: int = 0
|
total_len: int = 0
|
||||||
last_device_node: Any = None
|
last_device_node: Any = None
|
||||||
mamba_exist: bool = False
|
mamba_exist: bool = False
|
||||||
|
swa_branch_inserted: bool = False
|
||||||
inserted_host_node: Any = None
|
inserted_host_node: Any = None
|
||||||
host_insert_dropped: bool = False
|
host_insert_dropped: bool = False
|
||||||
adopted_ranges: Optional[dict[ComponentType, list[tuple[int, int]]]] = None
|
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;
|
loaded back to device. Pure-KV cache semantics;
|
||||||
swa_host_hit_length : Number of SWA tokens that hit on host (within the sliding
|
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.
|
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
|
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.
|
into the Mamba device pool. Typically 0 or 1.
|
||||||
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
|
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
|
||||||
@@ -216,6 +221,7 @@ class MatchResult(NamedTuple):
|
|||||||
best_match_node: Any
|
best_match_node: Any
|
||||||
host_hit_length: int = 0
|
host_hit_length: int = 0
|
||||||
swa_host_hit_length: int = 0
|
swa_host_hit_length: int = 0
|
||||||
|
swa_branching_seqlen: Optional[int] = None
|
||||||
mamba_host_hit_length: int = 0
|
mamba_host_hit_length: int = 0
|
||||||
mamba_branching_seqlen: Optional[int] = None
|
mamba_branching_seqlen: Optional[int] = None
|
||||||
cache_protected_len: Optional[int] = None
|
cache_protected_len: Optional[int] = None
|
||||||
@@ -240,6 +246,7 @@ def zero_match_result(
|
|||||||
best_match_node=root,
|
best_match_node=root,
|
||||||
host_hit_length=0,
|
host_hit_length=0,
|
||||||
swa_host_hit_length=0,
|
swa_host_hit_length=0,
|
||||||
|
swa_branching_seqlen=None,
|
||||||
mamba_host_hit_length=0,
|
mamba_host_hit_length=0,
|
||||||
full_kv_hit_length=0,
|
full_kv_hit_length=0,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Sequence
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||||
DecLockRefParams,
|
DecLockRefParams,
|
||||||
EvictParams,
|
EvictParams,
|
||||||
@@ -84,8 +85,32 @@ class SWAComponent(TreeComponent):
|
|||||||
|
|
||||||
component_type = ComponentType.SWA
|
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:
|
def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool:
|
||||||
return False
|
return bool(self._dirty_backup_window(node))
|
||||||
|
|
||||||
def reset_session_state(self) -> None:
|
def reset_session_state(self) -> None:
|
||||||
super().reset_session_state()
|
super().reset_session_state()
|
||||||
@@ -308,6 +333,16 @@ class SWAComponent(TreeComponent):
|
|||||||
best_value_len: int,
|
best_value_len: int,
|
||||||
) -> MatchResult:
|
) -> MatchResult:
|
||||||
ct = self.component_type
|
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
|
n_swa = 0
|
||||||
swa_host_hit = 0
|
swa_host_hit = 0
|
||||||
node = result.best_match_node
|
node = result.best_match_node
|
||||||
@@ -328,11 +363,11 @@ class SWAComponent(TreeComponent):
|
|||||||
else:
|
else:
|
||||||
break
|
break
|
||||||
node = node.parent
|
node = node.parent
|
||||||
if swa_host_hit > 0:
|
|
||||||
return result._replace(
|
return result._replace(
|
||||||
swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit)
|
swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit),
|
||||||
|
swa_branching_seqlen=branching_seqlen,
|
||||||
)
|
)
|
||||||
return result
|
|
||||||
|
|
||||||
def update_component_on_insert_overlap(
|
def update_component_on_insert_overlap(
|
||||||
self,
|
self,
|
||||||
@@ -466,6 +501,11 @@ class SWAComponent(TreeComponent):
|
|||||||
result: InsertResult,
|
result: InsertResult,
|
||||||
cache_actions: list[CacheAction | ComponentAction],
|
cache_actions: list[CacheAction | ComponentAction],
|
||||||
) -> None:
|
) -> 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:
|
if not is_new_leaf:
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -857,12 +897,23 @@ class SWAComponent(TreeComponent):
|
|||||||
# Unfinished requests can already have an SWA-evicted prefix; preserve
|
# Unfinished requests can already have an SWA-evicted prefix; preserve
|
||||||
# that boundary so insertion creates a tombstone instead of live SWA KV.
|
# that boundary so insertion creates a tombstone instead of live SWA KV.
|
||||||
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
|
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
|
||||||
|
|
||||||
|
branching_seqlen = req.swa_branching_seqlen
|
||||||
|
if branching_seqlen is None or branching_seqlen <= req.kv.cache_protected_len:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def free_out_of_window_slots(
|
# An EAGLE key with N bigrams spans N + 1 raw tokens.
|
||||||
self, req: Req, pre_len: int, insert_params: InsertParams
|
effective_cache_len = branching_seqlen + int(self.tree_core.is_eagle)
|
||||||
) -> None:
|
if effective_cache_len > token_ids_len:
|
||||||
if self.sliding_window_size is not None:
|
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(
|
free_swa_out_of_window_slots(
|
||||||
req,
|
req,
|
||||||
pre_len,
|
pre_len,
|
||||||
@@ -872,8 +923,33 @@ class SWAComponent(TreeComponent):
|
|||||||
token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
|
token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
|
||||||
retain_floor=self.cache.swa_retain_floor(req),
|
retain_floor=self.cache.swa_retain_floor(req),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def free_out_of_window_slots(
|
||||||
|
self, req: Req, pre_len: int, insert_params: InsertParams
|
||||||
|
) -> None:
|
||||||
|
self._free_out_of_window_slots(req, pre_len)
|
||||||
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
|
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 ----
|
# ---- HiCache Hooks ----
|
||||||
|
|
||||||
def prepare_prefetch(
|
def prepare_prefetch(
|
||||||
@@ -933,15 +1009,22 @@ class SWAComponent(TreeComponent):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
if phase == CacheTransferPhase.BACKUP_HOST:
|
if phase == CacheTransferPhase.BACKUP_HOST:
|
||||||
|
if self.cache.host_memory_mode == "buffer_only":
|
||||||
|
# Buffer mode stages one node/hash span per FIFO backup intent.
|
||||||
cd = node.component_data[ct]
|
cd = node.component_data[ct]
|
||||||
if cd.value is None:
|
dirty = [node] if cd.value is not None else []
|
||||||
|
else:
|
||||||
|
dirty = self._dirty_backup_window(node)
|
||||||
|
if not dirty:
|
||||||
return None
|
return None
|
||||||
# cd.value already holds SWA-pool indices (translated at insert time).
|
dirty.reverse()
|
||||||
# Host pool indexing wants int64.
|
|
||||||
return [
|
return [
|
||||||
PoolTransfer(
|
PoolTransfer(
|
||||||
name=PoolName.SWA,
|
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
|
ct = self.component_type
|
||||||
|
|
||||||
if phase == CacheTransferPhase.BACKUP_HOST:
|
if phase == CacheTransferPhase.BACKUP_HOST:
|
||||||
if transfers and transfers[0].host_indices is not None:
|
if not transfers or transfers[0].host_indices is None:
|
||||||
cd = node.component_data[ct]
|
return
|
||||||
if cd.host_value is None:
|
xfer = transfers[0]
|
||||||
cd.host_value = transfers[0].host_indices.clone()
|
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
|
return
|
||||||
|
|
||||||
if phase == CacheTransferPhase.LOAD_BACK:
|
if phase == CacheTransferPhase.LOAD_BACK:
|
||||||
|
|||||||
@@ -1131,9 +1131,13 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
|
|||||||
state.phase = _InsertPhase.TAIL
|
state.phase = _InsertPhase.TAIL
|
||||||
|
|
||||||
def _needs_incremental_component_backup(self, node: UnifiedTreeNode) -> bool:
|
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(
|
return any(
|
||||||
component.needs_incremental_backup(node)
|
component.needs_incremental_backup(node)
|
||||||
for component in self.components
|
for component in components
|
||||||
if component.component_type != BASE_COMPONENT_TYPE
|
if component.component_type != BASE_COMPONENT_TYPE
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1147,7 +1151,6 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
|
|||||||
node = state.target_node
|
node = state.target_node
|
||||||
return (
|
return (
|
||||||
self.enable_hicache
|
self.enable_hicache
|
||||||
and not self.is_write_back
|
|
||||||
and node.backuped
|
and node.backuped
|
||||||
and node.write_through_pending_id is None
|
and node.write_through_pending_id is None
|
||||||
and self._needs_incremental_component_backup(node)
|
and self._needs_incremental_component_backup(node)
|
||||||
@@ -1979,7 +1982,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
|
|||||||
for comp in self.components:
|
for comp in self.components:
|
||||||
if comp.component_type == BASE_COMPONENT_TYPE:
|
if comp.component_type == BASE_COMPONENT_TYPE:
|
||||||
continue
|
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
|
continue
|
||||||
t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST)
|
t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST)
|
||||||
if t:
|
if t:
|
||||||
|
|||||||
@@ -66,6 +66,8 @@ class SessionSlot:
|
|||||||
# Later turns run on the slot's record (see restore_to_req).
|
# Later turns run on the slot's record (see restore_to_req).
|
||||||
assert kv is self.kv
|
assert kv is self.kv
|
||||||
|
|
||||||
|
req.swa_branching_seqlen = None
|
||||||
|
|
||||||
def restore_to_req(self, req: Req):
|
def restore_to_req(self, req: Req):
|
||||||
"""Restore KV state from this slot into an incoming request."""
|
"""Restore KV state from this slot into an incoming request."""
|
||||||
req.kv = self.kv
|
req.kv = self.kv
|
||||||
|
|||||||
@@ -39,6 +39,7 @@ class _StubReq:
|
|||||||
self.best_match_node = None
|
self.best_match_node = None
|
||||||
self.host_hit_length = None
|
self.host_hit_length = None
|
||||||
self.num_matched_prefix_tokens = 0
|
self.num_matched_prefix_tokens = 0
|
||||||
|
self.swa_branching_seqlen = None
|
||||||
self.kv = SimpleNamespace(cache_protected_len=None)
|
self.kv = SimpleNamespace(cache_protected_len=None)
|
||||||
|
|
||||||
def _compute_max_prefix_len(self, input_len):
|
def _compute_max_prefix_len(self, input_len):
|
||||||
@@ -79,6 +80,28 @@ class TestZeroMatchResult(unittest.TestCase):
|
|||||||
|
|
||||||
|
|
||||||
class TestMatchPrefixForReqForceMiss(unittest.TestCase):
|
class TestMatchPrefixForReqForceMiss(unittest.TestCase):
|
||||||
|
def test_swa_branching_seqlen_is_cleared_without_new_branch(self):
|
||||||
|
class _StubTreeCache:
|
||||||
|
def swa_reprefill_tail_tokens(self):
|
||||||
|
return 0
|
||||||
|
|
||||||
|
def match_prefix(self, params):
|
||||||
|
return MatchResult(
|
||||||
|
device_indices=torch.empty((0,), dtype=torch.int64),
|
||||||
|
last_device_node=None,
|
||||||
|
last_host_node=None,
|
||||||
|
best_match_node=None,
|
||||||
|
host_hit_length=0,
|
||||||
|
swa_branching_seqlen=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
req = _StubReq([1, 2, 3, 4])
|
||||||
|
req.swa_branching_seqlen = 8
|
||||||
|
|
||||||
|
match_prefix_for_req(_StubTreeCache(), req)
|
||||||
|
|
||||||
|
self.assertIsNone(req.swa_branching_seqlen)
|
||||||
|
|
||||||
def test_force_miss_zeros_req_prefix(self):
|
def test_force_miss_zeros_req_prefix(self):
|
||||||
tree = RadixCache.create_simulated()
|
tree = RadixCache.create_simulated()
|
||||||
tree.insert(
|
tree.insert(
|
||||||
|
|||||||
@@ -99,6 +99,7 @@ class _FakeReq:
|
|||||||
self.last_node = None
|
self.last_node = None
|
||||||
self.swa_uuid_for_lock = None
|
self.swa_uuid_for_lock = None
|
||||||
self.skip_lock_node_ids = {}
|
self.skip_lock_node_ids = {}
|
||||||
|
self.swa_branching_seqlen = None
|
||||||
self.to_finish = None
|
self.to_finish = None
|
||||||
self.finished_reason = None
|
self.finished_reason = None
|
||||||
self.finished_len = None
|
self.finished_len = None
|
||||||
@@ -265,6 +266,20 @@ def test_release_session_threads_mamba_skip_ids():
|
|||||||
assert params.skip_lock_node_ids.get(ComponentType.MAMBA) == {42}
|
assert params.skip_lock_node_ids.get(ComponentType.MAMBA) == {42}
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_slot_does_not_restore_swa_branching_seqlen():
|
||||||
|
req = _FakeReq("session-a", req_pool_idx=0, committed=4, allocated=4)
|
||||||
|
req.swa_branching_seqlen = 8
|
||||||
|
|
||||||
|
slot = SessionSlot()
|
||||||
|
slot.save_from_req(req, is_first=True)
|
||||||
|
|
||||||
|
next_req = _FakeReq("session-a", req_pool_idx=1, committed=0, allocated=0)
|
||||||
|
slot.restore_to_req(next_req)
|
||||||
|
|
||||||
|
assert req.swa_branching_seqlen is None
|
||||||
|
assert next_req.swa_branching_seqlen is None
|
||||||
|
|
||||||
|
|
||||||
# Shrink tests removed: streaming sessions are append-only after the
|
# Shrink tests removed: streaming sessions are append-only after the
|
||||||
# rollback fix in session_controller (rollback_aborted_req). The shrink
|
# rollback fix in session_controller (rollback_aborted_req). The shrink
|
||||||
# code path in cache_finished_req no longer exists.
|
# code path in cache_finished_req no longer exists.
|
||||||
|
|||||||
@@ -5765,6 +5765,231 @@ class UnifiedRadixCacheSuite:
|
|||||||
self.assertEqual(result.host_hit_length, 0)
|
self.assertEqual(result.host_hit_length, 0)
|
||||||
self.assertEqual(result.swa_host_hit_length, _node_key_length(cache, leaf))
|
self.assertEqual(result.swa_host_hit_length, _node_key_length(cache, leaf))
|
||||||
|
|
||||||
|
def _skip_swa_branching_on_rust(self) -> None:
|
||||||
|
# TODO(alphabetc1): drop this gate once #37584 ports SWA branching-point
|
||||||
|
# caching to the Rust tree core.
|
||||||
|
if _selected_tree_core_test_backend() == "rust":
|
||||||
|
self.skipTest("SWA branching-point caching is Python-core only")
|
||||||
|
|
||||||
|
def test_swa_branching_seqlen_uses_device_full_hit(self):
|
||||||
|
self._skip_swa_branching_on_rust()
|
||||||
|
if (
|
||||||
|
not self.cfg.has_swa
|
||||||
|
or self.cfg.has_mamba
|
||||||
|
or self.cfg.page_size not in (1, 4)
|
||||||
|
or self.cfg.sliding_window_size != 4
|
||||||
|
):
|
||||||
|
self.skipTest("requires Full+SWA with window_size=4")
|
||||||
|
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||||
|
window = self.cfg.sliding_window_size
|
||||||
|
window_pages = (window + self.cfg.page_size - 1) // self.cfg.page_size
|
||||||
|
prefix = self._make_seq(1, window_pages)
|
||||||
|
tokens = prefix + self._make_seq(1000, 2)
|
||||||
|
self._insert(cache, allocator, req_to_token_pool, prefix)
|
||||||
|
self._insert(cache, allocator, req_to_token_pool, tokens)
|
||||||
|
|
||||||
|
leaf = cache.resolve_node_handle(
|
||||||
|
cache.match_prefix(
|
||||||
|
MatchPrefixParams(key=RadixKey(array("q", tokens)))
|
||||||
|
).last_device_node
|
||||||
|
)
|
||||||
|
device_frees = defaultdict(list)
|
||||||
|
cache.tree_core._evict_component_and_detach_lru(
|
||||||
|
leaf,
|
||||||
|
cache.components[ComponentType.SWA],
|
||||||
|
device_frees=device_frees,
|
||||||
|
host_frees=defaultdict(list),
|
||||||
|
target=EvictLayer.DEVICE,
|
||||||
|
)
|
||||||
|
cache._drain_device_frees(device_frees)
|
||||||
|
|
||||||
|
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
|
||||||
|
|
||||||
|
self.assertEqual(result.full_kv_hit_length, len(tokens))
|
||||||
|
self.assertEqual(result.swa_branching_seqlen, len(tokens))
|
||||||
|
self.assertEqual(result.swa_branching_seqlen % self.cfg.page_size, 0)
|
||||||
|
|
||||||
|
# Simulate forward producing fresh SWA KV at the branching point.
|
||||||
|
self._insert(
|
||||||
|
cache,
|
||||||
|
allocator,
|
||||||
|
req_to_token_pool,
|
||||||
|
tokens[: result.swa_branching_seqlen],
|
||||||
|
)
|
||||||
|
|
||||||
|
rematch = cache.match_prefix(
|
||||||
|
MatchPrefixParams(key=RadixKey(array("q", tokens)))
|
||||||
|
)
|
||||||
|
self.assertEqual(len(rematch.device_indices), result.swa_branching_seqlen)
|
||||||
|
self.assertIsNone(rematch.swa_branching_seqlen)
|
||||||
|
|
||||||
|
def test_swa_branching_seqlen_uses_host_full_hit(self):
|
||||||
|
self._skip_swa_branching_on_rust()
|
||||||
|
if (
|
||||||
|
not self.cfg.has_swa
|
||||||
|
or self.cfg.has_mamba
|
||||||
|
or self.cfg.page_size != 1
|
||||||
|
or self.cfg.sliding_window_size != 4
|
||||||
|
):
|
||||||
|
self.skipTest("requires page_size=1 Full+SWA with window_size=4")
|
||||||
|
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||||
|
window = self.cfg.sliding_window_size
|
||||||
|
prefix = self._make_seq(1, window)
|
||||||
|
tokens = prefix + self._make_seq(1000, window + 1)
|
||||||
|
self._insert(cache, allocator, req_to_token_pool, prefix)
|
||||||
|
self._insert(cache, allocator, req_to_token_pool, tokens)
|
||||||
|
|
||||||
|
leaf = cache.resolve_node_handle(
|
||||||
|
cache.match_prefix(
|
||||||
|
MatchPrefixParams(key=RadixKey(array("q", tokens)))
|
||||||
|
).last_device_node
|
||||||
|
)
|
||||||
|
parent = leaf.parent
|
||||||
|
self._backup_node(cache, leaf.id)
|
||||||
|
lock_result = cache.inc_lock_ref(parent.id)
|
||||||
|
try:
|
||||||
|
cache.evict(EvictParams(num_tokens=len(leaf.key)))
|
||||||
|
finally:
|
||||||
|
cache.dec_lock_ref(parent.id, lock_result.to_dec_params())
|
||||||
|
device_frees = defaultdict(list)
|
||||||
|
host_frees = defaultdict(list)
|
||||||
|
cache.components[ComponentType.SWA].evict_component(
|
||||||
|
leaf, device_frees, host_frees, target=EvictLayer.HOST
|
||||||
|
)
|
||||||
|
cache._free_values(device_frees, host_frees)
|
||||||
|
full_host_pool = cache.cache_controller.mem_pool_host
|
||||||
|
swa_host_pool = cache.components[ComponentType.SWA]._swa_kv_pool_host
|
||||||
|
full_available_before = full_host_pool.available_size()
|
||||||
|
swa_available_before = swa_host_pool.available_size()
|
||||||
|
|
||||||
|
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
|
||||||
|
|
||||||
|
self.assertEqual(result.full_kv_hit_length, len(tokens))
|
||||||
|
self.assertEqual(result.swa_branching_seqlen, len(tokens))
|
||||||
|
|
||||||
|
self._insert(
|
||||||
|
cache,
|
||||||
|
allocator,
|
||||||
|
req_to_token_pool,
|
||||||
|
tokens[: result.swa_branching_seqlen],
|
||||||
|
)
|
||||||
|
cache.writing_check(write_back=True)
|
||||||
|
# Full was already backed up, so only the SWA window is allocated.
|
||||||
|
self.assertEqual(full_host_pool.available_size(), full_available_before)
|
||||||
|
self.assertEqual(
|
||||||
|
swa_host_pool.available_size(),
|
||||||
|
swa_available_before - len(leaf.key),
|
||||||
|
)
|
||||||
|
|
||||||
|
rematch = cache.match_prefix(
|
||||||
|
MatchPrefixParams(key=RadixKey(array("q", tokens)))
|
||||||
|
)
|
||||||
|
self.assertEqual(len(rematch.device_indices), result.swa_branching_seqlen)
|
||||||
|
self.assertIsNone(rematch.swa_branching_seqlen)
|
||||||
|
|
||||||
|
def test_swa_branching_seqlen_caps_insert_after_forward(self):
|
||||||
|
if (
|
||||||
|
not self.cfg.has_swa
|
||||||
|
or self.cfg.has_mamba
|
||||||
|
or self.cfg.page_size != 1
|
||||||
|
or self.cfg.sliding_window_size != 4
|
||||||
|
):
|
||||||
|
self.skipTest("requires page_size=1 Full+SWA with window_size=4")
|
||||||
|
cache, _, _ = build_fixture(self.cfg)
|
||||||
|
swa = cache.components[ComponentType.SWA]
|
||||||
|
req = mock.Mock(
|
||||||
|
swa_branching_seqlen=8,
|
||||||
|
kv=mock.Mock(cache_protected_len=4, swa_evicted_seqlen=0),
|
||||||
|
)
|
||||||
|
|
||||||
|
for is_finished in (False, True):
|
||||||
|
params = InsertParams()
|
||||||
|
self.assertEqual(
|
||||||
|
swa.prepare_for_caching_req(req, params, 12, is_finished), 8
|
||||||
|
)
|
||||||
|
self.assertEqual(params.swa_evicted_seqlen, 0)
|
||||||
|
|
||||||
|
self.assertIsNone(swa.prepare_for_caching_req(req, InsertParams(), 7, False))
|
||||||
|
req.kv.cache_protected_len = 8
|
||||||
|
self.assertIsNone(swa.prepare_for_caching_req(req, InsertParams(), 12, False))
|
||||||
|
|
||||||
|
cache.tree_core.is_eagle = True
|
||||||
|
req.kv.cache_protected_len = 4
|
||||||
|
params = InsertParams()
|
||||||
|
self.assertEqual(swa.prepare_for_caching_req(req, params, 12, False), 9)
|
||||||
|
params.key = RadixKey(array("q", range(9)), is_bigram=True)
|
||||||
|
result = mock.Mock(swa_branch_inserted=False)
|
||||||
|
swa.commit_insert_component_data(mock.Mock(), False, params, result, [])
|
||||||
|
self.assertTrue(result.swa_branch_inserted)
|
||||||
|
|
||||||
|
def test_swa_branch_insert_releases_forward_overshoot(self):
|
||||||
|
if (
|
||||||
|
not self.cfg.has_swa
|
||||||
|
or self.cfg.has_mamba
|
||||||
|
or self.cfg.page_size != 1
|
||||||
|
or self.cfg.sliding_window_size != 4
|
||||||
|
):
|
||||||
|
self.skipTest("requires page_size=1 Full+SWA with window_size=4")
|
||||||
|
|
||||||
|
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||||
|
swa = cache.components[ComponentType.SWA]
|
||||||
|
|
||||||
|
branching_seqlen = 8
|
||||||
|
forward_len = 20
|
||||||
|
req = self._make_req(req_to_token_pool)
|
||||||
|
tokens = self._make_seq(1, forward_len)
|
||||||
|
req.origin_input_ids = tokens
|
||||||
|
req.output_ids = []
|
||||||
|
req.full_untruncated_fill_ids = array("q", tokens)
|
||||||
|
req.set_extend_range(0, forward_len)
|
||||||
|
req.kv.cache_protected_len = branching_seqlen
|
||||||
|
|
||||||
|
kv_indices = self._alloc(allocator, forward_len)
|
||||||
|
req_to_token_pool.write(
|
||||||
|
(req.kv.req_pool_idx, slice(0, forward_len)),
|
||||||
|
kv_indices,
|
||||||
|
)
|
||||||
|
|
||||||
|
params = InsertParams(
|
||||||
|
key=RadixKey(array("q", tokens[:branching_seqlen])),
|
||||||
|
swa_branching_seqlen=branching_seqlen,
|
||||||
|
)
|
||||||
|
result = mock.Mock(swa_branch_inserted=False)
|
||||||
|
swa.commit_insert_component_data(
|
||||||
|
mock.Mock(),
|
||||||
|
False,
|
||||||
|
params,
|
||||||
|
result,
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
self.assertTrue(result.swa_branch_inserted)
|
||||||
|
|
||||||
|
with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
|
||||||
|
swa.cleanup_after_caching_req(
|
||||||
|
req,
|
||||||
|
is_finished=False,
|
||||||
|
insert_result=result,
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_evicted = forward_len - 1 - self.cfg.sliding_window_size
|
||||||
|
self.assertEqual(req.kv.swa_evicted_seqlen, expected_evicted)
|
||||||
|
|
||||||
|
mapping = allocator.full_to_swa_index_mapping
|
||||||
|
self.assertEqual(
|
||||||
|
torch.count_nonzero(mapping[kv_indices[:branching_seqlen]]).item(),
|
||||||
|
branching_seqlen,
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
torch.count_nonzero(
|
||||||
|
mapping[kv_indices[branching_seqlen:expected_evicted]]
|
||||||
|
).item(),
|
||||||
|
0,
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
torch.count_nonzero(mapping[kv_indices[expected_evicted:]]).item(),
|
||||||
|
forward_len - expected_evicted,
|
||||||
|
)
|
||||||
|
|
||||||
def test_mamba_branching_seqlen_disabled_under_hicache(self):
|
def test_mamba_branching_seqlen_disabled_under_hicache(self):
|
||||||
if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
|
if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1:
|
||||||
self.skipTest("requires page_size=1 Full+Mamba")
|
self.skipTest("requires page_size=1 Full+Mamba")
|
||||||
@@ -6467,6 +6692,66 @@ class UnifiedRadixCacheSuite:
|
|||||||
self.assertEqual(kv_xfer.nodes_to_load, [b])
|
self.assertEqual(kv_xfer.nodes_to_load, [b])
|
||||||
self.assertEqual(comp_xfers[ComponentType.SWA][0].nodes_to_load, [a, b])
|
self.assertEqual(comp_xfers[ComponentType.SWA][0].nodes_to_load, [a, b])
|
||||||
|
|
||||||
|
def test_hicache_swa_backup_window_stops_at_pending_ancestor(self):
|
||||||
|
self._skip_swa_branching_on_rust()
|
||||||
|
if (
|
||||||
|
not self.cfg.has_swa
|
||||||
|
or self.cfg.has_mamba
|
||||||
|
or self.cfg.page_size != 1
|
||||||
|
or self.cfg.sliding_window_size != 4
|
||||||
|
):
|
||||||
|
self.skipTest("requires page_size=1 Full+SWA with window_size=4")
|
||||||
|
|
||||||
|
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||||
|
chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3)
|
||||||
|
if len(chain) < 3:
|
||||||
|
self.skipTest("chain collapsed below the pending ancestor test")
|
||||||
|
|
||||||
|
c = chain[-1]
|
||||||
|
c_swa = _device_value(cache, c, ComponentType.SWA).clone()
|
||||||
|
|
||||||
|
# First transfer: publish Full for C only, leaving SWA dirty while the
|
||||||
|
# write-through ack is still pending.
|
||||||
|
cache.tree_core.set_component_device_value_raw(c, ComponentType.SWA, None)
|
||||||
|
self.assertGreater(
|
||||||
|
cache._execute_and_commit_kv_backup(BackupKV(node_ids=[c])),
|
||||||
|
0,
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
cache.tree_core.node_by_id(c).write_through_pending_id,
|
||||||
|
c,
|
||||||
|
)
|
||||||
|
self.assertIsNotNone(_host_value(cache, c, ComponentType.FULL))
|
||||||
|
self.assertIsNone(_host_value(cache, c, ComponentType.SWA))
|
||||||
|
|
||||||
|
# Simulate SWA being reconstructed on device before the first ack. The
|
||||||
|
# next incremental SWA backup must treat C as the boundary and back up
|
||||||
|
# only the newly inserted descendant.
|
||||||
|
cache.tree_core.set_component_device_value_raw(c, ComponentType.SWA, c_swa)
|
||||||
|
tokens = self._match_tokens_for_chain(cache, chain)
|
||||||
|
next_tokens = tokens + self._make_seq(9000, 1)
|
||||||
|
cache.write_through_threshold = 1
|
||||||
|
|
||||||
|
self._insert(cache, allocator, req_to_token_pool, next_tokens)
|
||||||
|
|
||||||
|
d = cache.match_prefix(
|
||||||
|
MatchPrefixParams(key=RadixKey(array("q", next_tokens)))
|
||||||
|
).last_device_node
|
||||||
|
self.assertEqual(
|
||||||
|
cache.tree_core.node_by_id(c).write_through_pending_id,
|
||||||
|
c,
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
cache.tree_core.node_by_id(d).write_through_pending_id,
|
||||||
|
d,
|
||||||
|
)
|
||||||
|
self.assertIsNone(_host_value(cache, c, ComponentType.SWA))
|
||||||
|
self.assertIsNotNone(_host_value(cache, d, ComponentType.SWA))
|
||||||
|
|
||||||
|
cache.writing_check(write_back=True)
|
||||||
|
self.assertIsNone(cache.tree_core.node_by_id(c).write_through_pending_id)
|
||||||
|
self.assertIsNone(cache.tree_core.node_by_id(d).write_through_pending_id)
|
||||||
|
|
||||||
def _swa_finalize_setup(self):
|
def _swa_finalize_setup(self):
|
||||||
"""Build a SWA chain long enough to fill at least the window
|
"""Build a SWA chain long enough to fill at least the window
|
||||||
plus one extra page, and host-back every node so we can flip
|
plus one extra page, and host-back every node so we can flip
|
||||||
|
|||||||
Reference in New Issue
Block a user