Skip to content

Add DeepSeek-V4 decode radix cache with hierarchical reuse #39055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
huangyg-star wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add DeepSeek-V4 decode radix cache with hierarchical reuse #39055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 44 additions & 21 deletions python/sglang/srt/disaggregation/decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,18 @@ def _release_matched_prefix_lock(self, req: Req) -> None:
else:
self.tree_cache.dec_lock_ref(req.last_node, req.lock_receipt)

@staticmethod
def _has_matched_prefix_lock(
decode_req: DecodeRequest, prefix_match: Optional[DecodePrefixMatch]
) -> bool:
# A request restored from L2/L3 may have no device hit at all while its
# FULL admission lock is still held (only the SWA half was released
# early by the full-only lock convention).
return prefix_match is not None and (
prefix_match.l1_prefix_len > 0
or getattr(decode_req.req, "swa_prefix_lock_released", False)
)

def _reclaim_swa_tail_capacity(
self, swa_tail_len: int, req_id: str
) -> Optional[str]:
Expand Down Expand Up @@ -1223,23 +1235,29 @@ def pop_preallocated(
# Cap full-attention prefix reuse at the sliding-window start so
# the SWA window lands entirely in the fresh delta, keeping
# alloc_extend_swa_tail's tail->full mapping in range. Costs reuse
# of only the last ~window_size full-attention tokens.
if uses_swa_tail_prealloc and prefix_len > 0:
# of only the last ~window_size full-attention tokens. The cap is
# the page-aligned window start, so the trimmed restore stays
# page aligned.
if uses_swa_tail_prealloc and total_prefix_len > 0:
swa_prefix_cap = fill_len - self._swa_tail_len(fill_len)
if prefix_len > swa_prefix_cap:
prefix_len = swa_prefix_cap
prefix_indices = prefix_indices[:prefix_len]
# Cap the prefill-committed prefix too: tokens past the
# cap are not device-resident, so prefill must transfer
# them.
total_prefix_len = prefix_len
if prefix_match.decode_prefix_len > swa_prefix_cap:
# The whole match -- device slice first, then the L3/L2
# restore -- shrinks to the cap; tokens past it are
# re-transferred by prefill.
prefix_match.cap_restore(swa_prefix_cap)
prefix_len = prefix_match.l1_prefix_len
prefix_indices = prefix_match.prefix_indices
total_prefix_len = prefix_match.decode_prefix_len

# Decode transfers the SWA tail fresh, so retain only the
# full-attention prefix lock needed for reuse.
if (
uses_swa_tail_prealloc
and prefix_match.l1_prefix_len > 0
and hasattr(self.tree_cache, "dec_swa_lock_only")
# full-attention prefix lock needed for reuse. Unconditional for
# SWA-tail models: with L2/L3 a request can have no device hit
# but still restore a prefix, and the restored node's lock
# (_try_hicache_queue_load_back) follows this same convention.
# dec_swa_lock_only walks up from the node, so a root last_node
# (no hit at all) is a no-op.
if uses_swa_tail_prealloc and hasattr(
self.tree_cache, "dec_swa_lock_only"
):
self.tree_cache.dec_swa_lock_only(
decode_req.req.last_node, decode_req.req.lock_receipt
Expand Down Expand Up @@ -1288,11 +1306,11 @@ def pop_preallocated(
)
> full_allocatable_tokens
):
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
if self._has_matched_prefix_lock(decode_req, prefix_match):
self._release_matched_prefix_lock(decode_req.req)
break
if required_tokens_for_request > full_allocatable_tokens:
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
if self._has_matched_prefix_lock(decode_req, prefix_match):
self._release_matched_prefix_lock(decode_req.req)
break

Expand All @@ -1310,15 +1328,15 @@ def pop_preallocated(
)
> swa_allocatable_tokens
):
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
if self._has_matched_prefix_lock(decode_req, prefix_match):
self._release_matched_prefix_lock(decode_req.req)
break

reclaim_error = self._reclaim_swa_tail_capacity(
swa_len, decode_req.req.rid
)
if reclaim_error is not None:
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
if self._has_matched_prefix_lock(decode_req, prefix_match):
self._release_matched_prefix_lock(decode_req.req)
logger.error(reclaim_error)
prepare_abort(
Expand Down Expand Up @@ -1993,12 +2011,17 @@ def alloc_for_decode_prealloc(
)
if uses_swa_tail:
# Full-attention layers reuse prefix KV; SWA layers allocate only
# the live window tail.
# the live window tail. The full-attention prefix is the whole
# committed prefix: [prefix_len, total_prefix_len) is restored from
# host by HiCache, so only [total_prefix_len, fill_len) is allocated
# here. total_prefix_len is page aligned whenever it exceeds
# prefix_len, so last_loc (a page tail from the device prefix) is
# not consulted by alloc_extend.
kv_loc = allocator.alloc_extend_swa_tail(
prefix_lens=torch.tensor(
[prefix_len], dtype=torch.int64, device=device
[total_prefix_len], dtype=torch.int64, device=device
),
prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
prefix_lens_cpu=torch.tensor([total_prefix_len], dtype=torch.int64),
seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
last_loc=last_loc,
Expand Down
105 changes: 91 additions & 14 deletions python/sglang/srt/disaggregation/decode_hicache_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,18 @@ class DecodePrefixMatch:
last_device_node: Any
last_host_node: Any = None
prefetch_registered: bool = False
# The D node may prefetch the complete L3 candidate locally even when the
# prefix advertised to P is capped (for example by SWA-tail prealloc).
# ``l3_storage_hit_length`` remains the P-visible/restore contract.
raw_l3_storage_hit_length: Optional[int] = None
raw_l3_match_start: Optional[int] = None

def __post_init__(self) -> None:
# Keep direct construction of this dataclass backwards compatible.
if self.raw_l3_storage_hit_length is None:
self.raw_l3_storage_hit_length = self.l3_storage_hit_length
if self.raw_l3_match_start is None:
self.raw_l3_match_start = self.l1_prefix_len + self.l2_host_hit_length

@property
def l1_prefix_len(self) -> int:
Expand All @@ -48,6 +60,37 @@ def restore_token_count(self) -> int:
"""Number of tokens that need L2/L3 load_back to device."""
return self.decode_prefix_len - self.l1_prefix_len

def cap_restore(self, max_prefix_len: int) -> None:
"""Shrink the whole match so ``decode_prefix_len <= max_prefix_len``.

Used when the caller's layout only fits a shorter prefix (e.g. an
SWA-tail prealloc whose window must stay inside the fresh delta).
Trims the device-resident slice first, then the L3 and L2 tiers so the
surviving restore stays contiguous from the device prefix.
``max_prefix_len`` must be page aligned, otherwise the surviving
restore is not. Called before ``_start_hicache_prefetch``, so nothing
in the P-visible restore is kept for the trimmed part. The raw L3
candidate is intentionally unchanged: the D node still prefetches it
locally, even when the P-visible contract is capped.
"""
if self.l1_prefix_len > max_prefix_len:
self.prefix_indices = self.prefix_indices[:max_prefix_len]
excess = self.decode_prefix_len - max_prefix_len
if excess <= 0:
return
trimmed = min(self.l3_storage_hit_length, excess)
self.l3_storage_hit_length -= trimmed
excess -= trimmed
if excess > 0:
self.l2_host_hit_length -= min(self.l2_host_hit_length, excess)
if (
self.l3_storage_hit_length <= 0
and self.l2_host_hit_length <= 0
and self.raw_l3_storage_hit_length <= 0
):
# No host restore remains, so detach the host anchor.
self.last_host_node = None


class HiCacheRestoreResult(Enum):
"""Outcome of one tick of the HiCache local-restore state machine."""
Expand Down Expand Up @@ -98,8 +141,12 @@ def _build_decode_prefix_match(self, req: Req, result: Any) -> DecodePrefixMatch
l3_storage_hit_length=l3_storage_hit_length,
last_device_node=result.last_device_node,
last_host_node=(
result.last_host_node if l3_storage_hit_length > 0 else None
result.last_host_node
if l2_host_hit_length > 0 or l3_storage_hit_length > 0
else None
),
raw_l3_storage_hit_length=l3_storage_hit_length,
raw_l3_match_start=l1_prefix_len + l2_host_hit_length,
)

def _start_hicache_prefetch(
Expand All @@ -111,15 +158,14 @@ def _start_hicache_prefetch(
"""
if (
prefix_match is None
or prefix_match.l3_storage_hit_length <= 0
or prefix_match.raw_l3_storage_hit_length <= 0
or prefix_match.last_host_node is None
):
return
try:
matched_len = prefix_match.l1_prefix_len + prefix_match.l2_host_hit_length
suffix = req.origin_input_ids[
matched_len : matched_len + prefix_match.l3_storage_hit_length
]
matched_len = prefix_match.raw_l3_match_start
prefetch_length = prefix_match.raw_l3_storage_hit_length
suffix = req.origin_input_ids[matched_len : matched_len + prefetch_length]
last_hash = self.tree_cache.get_last_hash_value(prefix_match.last_host_node)
prefix_keys = (
self.tree_cache.get_prefix_hash_values(prefix_match.last_host_node)
Expand All @@ -140,7 +186,8 @@ def _start_hicache_prefetch(
)
except Exception as e:
logger.warning(
"HiCache L3 prefetch failed for rid=%s: %s; falling back to L2-only LoadingBack",
"HiCache L3 prefetch failed for rid=%s: %s; falling back to "
"L2-only LoadingBack",
req.rid,
e,
)
Expand Down Expand Up @@ -189,6 +236,7 @@ def _clean_hicache_prefetch_resources(self, decode_req: DecodeRequest) -> None:
self.tree_cache.dec_lock_ref(
decode_req.hicache_restored_node,
decode_req.hicache_restore_lock_receipt,
skip_swa=decode_req.req.swa_prefix_lock_released,
)
decode_req.hicache_restored_node = None
decode_req.hicache_restore_lock_receipt = None
Expand All @@ -203,13 +251,21 @@ def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool:
"""
pm = dr.prefix_match

# Wait for L3 -> L2 prefetch to drain (skip when no L3 hit).
if pm.l3_storage_hit_length > 0:
# Wait for an L3 contract or registered D-local prefetch to drain.
if pm.l3_storage_hit_length > 0 or pm.prefetch_registered:
if not self.tree_cache.check_prefetch_progress(dr.req.rid):
return False
self.tree_cache.pop_prefetch_loaded_tokens(dr.req.rid)

# Re-match: req.last_node / prefix_indices updated to current device state.
# A capped match can have no P-visible host restore while still
# carrying a D-local L3 prefetch. Drain that operation, but do not
# create a load_back request for the extra locally prefetched span.
if not pm.needs_local_restore:
dr.hicache_restore_status = HiCacheRestoreResult.READY
return False

# Re-match updates request metadata, but the admission node remains the
# request's lock owner until restore commit.
rematch = match_prefix_for_req(
self.tree_cache,
dr.req,
Expand Down Expand Up @@ -249,10 +305,29 @@ def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool:
dr.hicache_restored_kv_indices = torch.cat(
[rematch.device_indices[pm.l1_prefix_len :], new_indices]
)
if len(dr.hicache_restored_kv_indices) > pm.restore_token_count:
# The tree can grow between match and rematch (more pages backed up
# or a deeper match inserted), so the restored segment may cover
# more than the prefix promised to prefill. Only the promised part
# belongs to this request: writing past it would clobber the
# freshly preallocated delta slots in req_to_token.
dr.hicache_restored_kv_indices = dr.hicache_restored_kv_indices[
: pm.restore_token_count
]
dr.hicache_restored_node = restored_node
dr.hicache_restore_lock_receipt = self.tree_cache.inc_lock_ref(
restored_node
).to_dec_params()
lock_result = self.tree_cache.inc_lock_ref(restored_node)
dr.hicache_restore_lock_receipt = lock_result.to_dec_params()
if dr.req.swa_prefix_lock_released and hasattr(
self.tree_cache, "dec_swa_lock_only"
):
# The admission match already dropped the SWA half of its lock
# (decode transfers the SWA tail fresh), so keep the restored-node
# lock full-only too -- acquire then release, so the per-node
# lock_ref bookkeeping stays balanced.
self.tree_cache.dec_swa_lock_only(
restored_node,
dr.hicache_restore_lock_receipt,
)

if len(new_indices) == 0:
# Whole prefix already on device; no DMA needed.
Expand All @@ -271,7 +346,9 @@ def _process_hicache_local_restores(self, decode_reqs: List[DecodeRequest]) -> N
if dr.hicache_restore_status != HiCacheRestoreResult.PENDING:
continue
pm = dr.prefix_match
if pm is None or not pm.needs_local_restore:
if pm is None or (
not pm.needs_local_restore and not pm.prefetch_registered
):
dr.hicache_restore_status = HiCacheRestoreResult.READY
continue
active.append(dr)
Expand Down
7 changes: 7 additions & 0 deletions python/sglang/srt/mem_cache/base_prefix_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,13 @@ class InsertParams:
# off). See UnifiedTreeNode.rotation_base.
rotation_base: Optional[int] = None

# Optional fast path for a trusted request admission match whose device
# prefix remains locked and resident. UnifiedTreeCore revalidates the
# structural and component-residency invariants before using the anchor.
insert_anchor_node: Optional[Any] = None
insert_anchor_prefix_len: int = 0
insert_anchor_rid: Optional[Any] = None


@dataclasses.dataclass
class InsertResult:
Expand Down
26 changes: 20 additions & 6 deletions python/sglang/srt/mem_cache/kv_cache_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -257,24 +257,38 @@ def build_kv_cache(
and get_disagg().disaggregation_mode == "decode"
):
if is_hybrid_swa:
is_dsv4 = getattr(model_config, "is_deepseek_v4_arch", False)
if not (envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx()):
raise ValueError(
"--disaggregation-decode-enable-radix-cache with sliding "
"window attention (SWA) models requires the unified radix "
"tree (set SGLANG_ENABLE_UNIFIED_RADIX_TREE=1)."
)
if enable_hierarchical_cache:
if enable_hierarchical_cache and not is_dsv4:
raise ValueError(
"--disaggregation-decode-enable-radix-cache with sliding "
"window attention (SWA) models currently supports only "
"device-resident cache and is incompatible with "
"--enable-hierarchical-cache."
)
if getattr(model_config, "is_deepseek_v4_arch", False):
raise ValueError(
"--disaggregation-decode-enable-radix-cache does not support "
"DeepSeek-V4 (DSA) compressed KV (c4/c128/indexer) yet."
)
if is_dsv4:
# Reusing a prefix hands the matching compressed rows along with
# the full pages, and a compressed row only exists for a whole
# block. A page-aligned prefix is therefore block-aligned only
# while the page covers whole blocks -- which the DSV4 page-size
# override (256, or 128 on NPU) guarantees today.
compress_ratios = [
ratio
for ratio in getattr(model_config, "compress_ratios", None) or []
if ratio > 0
]
if compress_ratios and page_size % max(compress_ratios) != 0:
raise ValueError(
"--disaggregation-decode-enable-radix-cache with DeepSeek-V4 "
f"requires page_size ({page_size}) to be a multiple of the "
f"largest compression ratio ({max(compress_ratios)}) so that "
"a page-aligned prefix reuses whole compressed blocks."
)
if getattr(model_config, "is_hybrid_swa_compress", False):
raise ValueError(
"--disaggregation-decode-enable-radix-cache does not support "
Expand Down
Loading
Loading