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

[HiCache] Allow a retraction host pool smaller than the device pool (#35543)

Browse Source 
Co-authored-by: cctry <cctry@fb.com>
This commit is contained in:
cctry
2026-08-19 22:59:57 -07:00
committed by GitHub
co-authored by cctry
parent 50dae2d99d
commit 32d98aad13
7 changed files with 203 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+23 -7
View File
@@ -1914,7 +1914,8 @@ def release_req(
tree_cache: BasePrefixCache,
hisparse_coordinator: Optional[HiSparseCoordinator],
offload_kv: bool = True,
) -> None:
) -> bool:
"""Returns False when the KV backup failed and the request cannot be resumed."""
if hisparse_coordinator is not None and not req.finished():
hisparse_coordinator.retract_req(req)
@@ -1922,8 +1923,9 @@ def release_req(
# restored later without recompute (see resume_retracted_reqs/load_kv_cache).
# Callers that will recompute the KV instead (PD true-retraction rebootstrap)
# pass offload_kv=False to skip the wasteful device->host copy.
backup_saved = True
if server_args.disaggregation_mode == "decode" and offload_kv:
retraction_backup(
backup_saved = retraction_backup(
req,
tree_cache,
req_to_token_pool,
@@ -1937,6 +1939,7 @@ def release_req(
evict_from_tree_cache(tree_cache, num_tokens)
req.reset_for_retract()
return backup_saved
def retract_all(
@@ -2820,6 +2823,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
sorted_indices = self._get_decode_retraction_order(self.reqs, server_args)
retracted_reqs = []
reqs_to_abort: List[Req] = []
first_iter = True
while first_iter or (
not self.check_decode_mem(selected_indices=sorted_indices)
@@ -2831,11 +2835,23 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
first_iter = False
idx = sorted_indices.pop()
req = self.reqs[idx]
retracted_reqs.append(req)
# release memory and don't insert into the tree because we need the space instantly
self.release_req(idx, len(sorted_indices), server_args)
if self.release_req(idx, len(sorted_indices), server_args):
retracted_reqs.append(req)
else:
# The retraction host pool could not hold the backup and the
# device KV is already freed, so the request cannot resume.
req.to_finish = FINISH_ABORT(
"Retraction host KV pool exhausted. Aborting the request.",
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
)
reqs_to_abort.append(req)
logger.warning(
"retract_decode: aborted request %s, retraction host pool "
"exhausted",
req.rid,
)
reqs_to_abort: List[Req] = []
if len(sorted_indices) <= 1 and not self.check_decode_mem(
selected_indices=sorted_indices
):
@@ -2910,8 +2926,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
remaing_req_count: int,
server_args: ServerArgs,
offload_kv: bool = True,
):
release_req(
) -> bool:
return release_req(
req=self.reqs[idx],
remaing_req_count=remaing_req_count,
server_args=server_args,
python/sglang/srt/mem_cache/common.py
+6 -3
View File
@@ -144,17 +144,20 @@ def retraction_backup(
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
backend: str,
) -> None:
) -> bool:
"""Returns False when the host pool cannot hold the backup; the caller
aborts the request since its KV cannot be preserved."""
if backend == "cpu_tensor":
req.offload_kv_cache(req_to_token_pool, token_to_kv_pool_allocator)
return
return True
if backend != "host_pool":
raise ValueError(f"Unknown retraction backup backend: {backend}")
if req.seqlen <= 1:
return
return True
unified_cache = cast("UnifiedRadixCache", tree_cache)
req.retraction_backup = unified_cache.retraction_backup(req)
return req.retraction_backup is not None
def retraction_restore(
python/sglang/srt/mem_cache/kv_cache_builder.py
+14 -4
View File
@@ -139,6 +139,12 @@ def _register_legacy_hicache_draft(
tree_cache.cache_controller.set_draft_kv_pool(pool, draft_host_pool)
# Host slots a backup-only retraction pool gets, as a fraction of the device
# pool. Sized well under 1.0 because a retraction burst touches a fraction of
# the device tokens; overflow aborts the request rather than pre-reserving.
BACKUP_ONLY_HICACHE_RATIO = 0.2
def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
"""Resolve the retraction backend onto the config bags and return it.
@@ -180,10 +186,14 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
fields["disaggregation_decode_retraction_backup"] = backend
if memory.hicache_ratio is None:
# Only a decode server reaches resolution with the ratio unset; host-pool
# retraction sizes the host pool 1:1 with the device pool, everything
# else keeps the standard default.
fields["hicache_ratio"] = 1.0 if backend == "host_pool" else 2.0
# Only a decode server reaches resolution with the ratio unset. A
# backup-only pool can be small: retractions that overflow it abort their
# request instead of crashing the scheduler. Sharing the pool with
# HiCache keeps the standard default.
if backend == "host_pool" and not memory.enable_hierarchical_cache:
fields["hicache_ratio"] = BACKUP_ONLY_HICACHE_RATIO
else:
fields["hicache_ratio"] = 2.0
source = "kv_cache_builder.decode_retraction"
get_context().override(source, **fields)
python/sglang/srt/mem_cache/unified_radix_cache.py
+4 -28
View File
@@ -1042,24 +1042,6 @@ class UnifiedRadixCache(BasePrefixCache):
"an MHA or hybrid-SWA HiCache host stack."
)
kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
device_pools = {PoolName.KV: kv_cache}
if isinstance(kv_cache, SWAKVPool):
device_pools = {
PoolName.KV: kv_cache.full_kv_pool,
PoolName.SWA: kv_cache.swa_kv_pool,
}
for name, device_pool in device_pools.items():
host_pool = self.host_pool_group.entry_map[name].host_pool
if host_pool.logical_size < device_pool.size:
raise ValueError(
"Retraction host pool is smaller than its device pool: "
f"pool={name}, host_slots={host_pool.logical_size}, "
f"device_slots={device_pool.size}. Increase --hicache-ratio "
"or --hicache-size."
)
for spec in self.sidecar_pool_specs:
source_size = self.host_pool_group.entry_map[
spec.indices_from_pool
@@ -1138,7 +1120,8 @@ class UnifiedRadixCache(BasePrefixCache):
return 0
return self.evict_host(num_tokens)
def retraction_backup(self, req: Req) -> RetractionBackup:
def retraction_backup(self, req: Req) -> Optional[RetractionBackup]:
"""Back up device KV to the host pool; None when it cannot fit after reclaim."""
assert req.seqlen > 1
device_indices, extra_transfers = self._retraction_device_transfers(req)
@@ -1147,11 +1130,7 @@ class UnifiedRadixCache(BasePrefixCache):
self._reclaim_retraction_host(len(device_indices))
host_indices = self.host_pool_group.alloc(len(device_indices))
if host_indices is None:
raise RuntimeError(
"Retraction host KV pool exhausted after reclaim: "
f"request={req.rid}, required_slots={len(device_indices)}, "
f"available_slots={self.host_pool_group.available_size()}."
)
return None
resolved = self.cache_controller._resolve_pool_transfers_allocation(
extra_transfers or None,
@@ -1161,10 +1140,7 @@ class UnifiedRadixCache(BasePrefixCache):
)
if resolved is None and extra_transfers:
self.host_pool_group.free(host_indices)
raise RuntimeError(
"Retraction auxiliary host allocation failed after atomic rollback: "
f"request={req.rid}, pools={[x.name for x in extra_transfers]}."
)
return None
backup = RetractionBackup(
host_indices=host_indices,
python/sglang/srt/server_args.py
+1 -1
View File
@@ -2701,7 +2701,7 @@ class ServerArgs:
] = "cache"
hicache_ratio: A[
Optional[float],
"The ratio of the size of host KV cache memory pool to the size of device pool. Defaults to 2.0 in cache mode, 1.2 in buffer_only mode, or 1.0 for host-pool decode retraction.",
"The ratio of the size of host KV cache memory pool to the size of device pool. Defaults to 2.0 in cache mode, 1.2 in buffer_only mode, or 0.2 for backup-only host-pool decode retraction.",
NS("memory"),
] = None
hicache_size: A[