[AMD] Restore SWA reprefill-tail on UnifiedRadixCache when HiCache is off (#32759)
Co-authored-by: Thomas Wang <thomawan@amd.com>
This commit is contained in:
co-authored by
Thomas Wang
parent
570087ceda
commit
e9e9e37ddc
@@ -3071,20 +3071,28 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
|
|
||||||
def swa_reprefill_tail_tokens(self) -> int:
|
def swa_reprefill_tail_tokens(self) -> int:
|
||||||
"""
|
"""
|
||||||
Only unified_kv + HiCache needs this: SWA lives in a per-request ring
|
Only unified_kv needs this: SWA lives in a per-request ring
|
||||||
(state_slot/pos), not content-stable and never offloaded to host, so a
|
(state_slot/pos), not content-stable and never stored in the tree, so a
|
||||||
reused prefix's trailing sliding window would read another request's
|
reused prefix's trailing sliding window would read another request's
|
||||||
stale ring slots. Re-prefilling that window rewrites this request's ring
|
stale ring slots. Re-prefilling that window rewrites this request's ring.
|
||||||
(what plain radix reuse does via its SWA match gate). 0 for every other
|
|
||||||
layout.
|
Applies to plain radix reuse as well as HiCache -- the ring is stale
|
||||||
|
either way. Returns 0 once SWA has a host pool to restore exact contents
|
||||||
|
from, and for every non-unified_kv layout, whose SWA slots are
|
||||||
|
content-stable.
|
||||||
"""
|
"""
|
||||||
swa = self.components.get(ComponentType.SWA)
|
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||||
unified_compress_only_hicache = (
|
is_unified_kv_triton,
|
||||||
self.cache_controller is not None
|
|
||||||
and swa is not None
|
|
||||||
and not self.tree_core.has_swa_host_pool
|
|
||||||
)
|
)
|
||||||
return swa.sliding_window_size if unified_compress_only_hicache else 0
|
|
||||||
|
swa = self.components.get(ComponentType.SWA)
|
||||||
|
if swa is None or not swa.sliding_window_size:
|
||||||
|
return 0
|
||||||
|
if not is_unified_kv_triton():
|
||||||
|
return 0
|
||||||
|
if self.tree_core.has_swa_host_pool:
|
||||||
|
return 0
|
||||||
|
return swa.sliding_window_size
|
||||||
|
|
||||||
def swa_retain_floor(self, req) -> int | None:
|
def swa_retain_floor(self, req) -> int | None:
|
||||||
if not self.is_mamba_enabled or self._sliding_window_size is None:
|
if not self.is_mamba_enabled or self._sliding_window_size is None:
|
||||||
|
|||||||
Reference in New Issue
Block a user