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

[HiCache] Merge HiCache event checks to reduce decode overhead (#30511)

Browse Source 
Co-authored-by: alphabetc1 <47200617+alphabetc1@users.noreply.github.com>
This commit is contained in:
huangtingwei
2026-07-29 23:48:48 +08:00
committed by GitHub
co-authored by alphabetc1
parent bfc450248e
commit 50029f05a3
8 changed files with 219 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
-5
View File
@@ -3226,11 +3226,6 @@ class Scheduler(
batch.batch_is_full = False
return batch
# Eagerly release lock_ref on completed write-through nodes so they
# become evictable, improving batch scheduling headroom.
if self.enable_hierarchical_cache:
self.tree_cache.flush_write_through_acks()
# Check if decode out of memory
if (kv_full_retract_flag := not batch.check_decode_mem()) or (
TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0
python/sglang/srt/mem_cache/base_prefix_cache.py
-8
View File
@@ -364,14 +364,6 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
"""
raise NotImplementedError()
def flush_write_through_acks(self) -> None:
"""Release lock_ref on radix-tree nodes whose write-through has completed.
Lightweight operation that only processes finished write acks.
No-op for caches without hierarchical write-through support.
"""
pass
def check_hicache_events(self) -> Any:
"""
Check HiCache related activities to update radix tree and synchronize across TP workers if needed
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
-3
View File
@@ -466,9 +466,6 @@ class HiMambaRadixCache(MambaRadixCache):
def ready_to_load_host_cache(self) -> int:
return self.cache_controller.start_loading()
def flush_write_through_acks(self) -> None:
self.writing_check()
def check_hicache_events(self):
self.writing_check()
self.loading_check()
python/sglang/srt/mem_cache/hiradix_cache.py
+87 -32
View File
@@ -202,7 +202,6 @@ class HiRadixCache(RadixCache):
1 if server_args.hicache_write_policy == "write_through" else 2
)
self.load_back_threshold = 10
# Detach storage backend automatically on process shutdown
atexit.register(self.shutdown)
@@ -981,7 +980,42 @@ class HiRadixCache(RadixCache):
# write to host if the node is not backuped
self.write_backup(node)
def writing_check(self, write_back=False):
def _count_ready_acks(self, ack_queue) -> int:
ready_count = 0
for ack in ack_queue:
if not ack.finish_event.query():
break
ready_count += 1
return ready_count
def _sync_hicache_ready_counts(self) -> tuple[int, int, tuple[int, ...]]:
cache_controller = self.cache_controller
storage_queue_sizes = (
(
cache_controller.prefetch_revoke_queue.qsize(),
cache_controller.prefetch_hit_queue.qsize(),
cache_controller.ack_backup_queue.qsize(),
cache_controller.host_mem_release_queue.qsize(),
)
if self.enable_storage
else ()
)
ready_counts = torch.tensor(
[
self._count_ready_acks(cache_controller.ack_write_queue),
self._count_ready_acks(cache_controller.ack_load_queue),
*storage_queue_sizes,
],
dtype=torch.int,
device="cpu",
)
self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN)
count_values = list(map(int, ready_counts.tolist()))
return count_values[0], count_values[1], tuple(count_values[2:])
def writing_check(self, write_back=False, finish_count: Optional[int] = None):
if write_back:
# blocking till all write back complete
while len(self.ongoing_write_through) > 0:
@@ -993,19 +1027,21 @@ class HiRadixCache(RadixCache):
assert len(self.ongoing_write_through) == 0
return
# Every rank must enter the all_reduce below; ongoing_write_through can
# diverge across ranks (e.g. write_backup returning 0 on a subset under
# host memory pressure), so a conditional skip desyncs the NCCL op
# sequence and deadlocks under TP > 1. (Matches UnifiedRadixCache.)
finish_count = 0
if self.pp_rank == 0:
for ack in self.cache_controller.ack_write_queue:
if not ack.finish_event.query():
break
finish_count += 1
finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
if finish_count is None:
# Every rank must enter the all_reduce below; ongoing_write_through can
# diverge across ranks (e.g. write_backup returning 0 on a subset under
# host memory pressure), so a conditional skip desyncs the NCCL op
# sequence and deadlocks under TP > 1. (Matches UnifiedRadixCache.)
finish_count = 0
if self.pp_rank == 0:
finish_count = self._count_ready_acks(
self.cache_controller.ack_write_queue
)
finish_count_tensor = torch.tensor(
finish_count, dtype=torch.int, device="cpu"
)
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
if finish_count > 0:
logger.debug(f"Process {finish_count} write back operations")
@@ -1016,16 +1052,18 @@ class HiRadixCache(RadixCache):
self._finish_write_through_ack(ack_id, release_lock=True)
finish_count -= 1
def loading_check(self):
finish_count = 0
if self.pp_rank == 0:
for ack in self.cache_controller.ack_load_queue:
if not ack.finish_event.query():
break
finish_count += 1
finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
def loading_check(self, finish_count: Optional[int] = None):
if finish_count is None:
finish_count = 0
if self.pp_rank == 0:
finish_count = self._count_ready_acks(
self.cache_controller.ack_load_queue
)
finish_count_tensor = torch.tensor(
finish_count, dtype=torch.int, device="cpu"
)
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
if finish_count > 0:
logger.debug(f"Process {finish_count} load operations")
@@ -1440,16 +1478,33 @@ class HiRadixCache(RadixCache):
"""
return self.cache_controller.start_loading()
def flush_write_through_acks(self) -> None:
self.writing_check()
def check_hicache_events(self):
# Reap the previous round's PP-sync sends before issuing new ones.
self._drain_async_work()
self.writing_check()
self.loading_check()
if self.enable_storage:
self.drain_storage_control_queues()
if self.pp_size != 1:
self.writing_check()
self.loading_check()
if self.enable_storage:
self.drain_storage_control_queues()
else:
(
write_finish_count,
load_finish_count,
storage_queue_sizes,
) = self._sync_hicache_ready_counts()
self.writing_check(finish_count=write_finish_count)
self.loading_check(finish_count=load_finish_count)
if self.enable_storage and storage_queue_sizes:
n_revoke, n_storage_hit, n_backup, n_release = storage_queue_sizes[:4]
self._drain_storage_control_queues_impl(
n_revoke=n_revoke,
n_storage_hit=n_storage_hit,
n_backup=n_backup,
n_release=n_release,
log_metrics=True,
)
if self.enable_storage_metrics:
self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats()
python/sglang/srt/mem_cache/storage/flexkv/README.md
+2 -2
View File
@@ -379,8 +379,8 @@ Supported:
here but requires `PoolTransfer` + `PoolHitPolicy` plumbing in
`FlexKVConnector`.
* Write-back acks are per-request (one `dec_lock_ref` per
`cache_finished_req`), not per-page like HiCache's
`flush_write_through_acks`.
`cache_finished_req`), not per-page like HiCache's write-through
ack queues.
* `--radix-cache-backend=flexkv` and `--enable-flexkv` are
mutually equivalent today; we don't yet emit a deprecation
warning if both are set.
python/sglang/srt/mem_cache/unified_radix_cache.py
+114 -33
View File
@@ -1746,7 +1746,64 @@ class UnifiedRadixCache(BasePrefixCache):
# ---- HiCache: Async Event Management ----
def writing_check(self, write_back: bool = False) -> None:
def _count_ready_acks(self, ack_queue) -> int:
ready_count = 0
for ack in ack_queue:
if not ack.finish_event.query():
break
ready_count += 1
return ready_count
def _sync_hicache_ready_counts(
self,
) -> tuple[int, int, tuple[int, ...], tuple[PoolName, ...]]:
cc = self.cache_controller
if cc is None:
write_acks = 0
load_acks = 0
storage_queue_sizes = ()
extra_pool_names = ()
else:
write_acks = self._count_ready_acks(cc.ack_write_queue)
load_acks = self._count_ready_acks(cc.ack_load_queue)
extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {})
extra_pool_names = (
tuple(extra_release_queues) if self.enable_storage else ()
)
storage_queue_sizes = (
(
cc.prefetch_revoke_queue.qsize(),
cc.prefetch_hit_queue.qsize(),
cc.ack_backup_queue.qsize(),
cc.host_mem_release_queue.qsize(),
*(extra_release_queues[name].qsize() for name in extra_pool_names),
)
if self.enable_storage
else ()
)
ready_counts = torch.tensor(
[
write_acks,
load_acks,
*storage_queue_sizes,
],
dtype=torch.int,
device="cpu",
)
self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN)
count_values = list(map(int, ready_counts.tolist()))
return (
count_values[0],
count_values[1],
tuple(count_values[2:]),
extra_pool_names,
)
def writing_check(
self, write_back: bool = False, finish_count: Optional[int] = None
) -> None:
"""Poll write-through completions."""
cc = self.cache_controller
if cc is None:
@@ -1764,18 +1821,17 @@ class UnifiedRadixCache(BasePrefixCache):
assert len(self.ongoing_write_through) == 0
return
# Every rank must enter the all_reduce below; ongoing_write_through can
# diverge across ranks (e.g. a backup returning 0 on a subset).
finish_count = 0
if self.pp_rank == 0:
for ack in cc.ack_write_queue:
if not ack.finish_event.query():
break
finish_count += 1
finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
if finish_count is None:
# Every rank must enter the all_reduce below; ongoing_write_through can
# diverge across ranks (e.g. write_backup returning 0 on a subset).
finish_count = 0
if self.pp_rank == 0:
finish_count = self._count_ready_acks(cc.ack_write_queue)
finish_count_tensor = torch.tensor(
finish_count, dtype=torch.int, device="cpu"
)
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
# Process completed acks
while finish_count > 0:
@@ -1785,22 +1841,22 @@ class UnifiedRadixCache(BasePrefixCache):
self._finish_write_through_ack(ack_id)
finish_count -= 1
def loading_check(self) -> None:
def loading_check(self, finish_count: Optional[int] = None) -> None:
"""Poll load-back completions."""
cc = self.cache_controller
if cc is None:
return
# Every rank must enter the all_reduce below; ongoing_load_back can
# diverge across ranks.
finish_count = 0
if self.pp_rank == 0:
for ack in cc.ack_load_queue:
if not ack.finish_event.query():
break
finish_count += 1
finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu")
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
if finish_count is None:
# Every rank must enter the all_reduce below; ongoing_load_back can
# diverge across ranks.
finish_count = 0
if self.pp_rank == 0:
finish_count = self._count_ready_acks(cc.ack_load_queue)
finish_count_tensor = torch.tensor(
finish_count, dtype=torch.int, device="cpu"
)
self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN)
finish_count = finish_count_tensor.item()
while finish_count > 0:
ack = cc.ack_load_queue.pop(0)
@@ -1867,19 +1923,44 @@ class UnifiedRadixCache(BasePrefixCache):
"""Called per scheduler step to poll async HiCache events."""
# Reap the previous round's PP-sync sends before issuing new ones.
self._drain_async_work()
self.writing_check()
self.loading_check()
if self.enable_storage:
self.drain_storage_control_queues()
if self.pp_size != 1:
self.writing_check()
self.loading_check()
if self.enable_storage:
self.drain_storage_control_queues()
else:
(
write_finish_count,
load_finish_count,
storage_queue_sizes,
extra_pool_names,
) = self._sync_hicache_ready_counts()
self.writing_check(finish_count=write_finish_count)
self.loading_check(finish_count=load_finish_count)
if self.enable_storage and storage_queue_sizes:
n_revoke, n_storage_hit, n_backup, n_release = storage_queue_sizes[:4]
extra_release_counts = {
pool_name: count
for pool_name, count in zip(
extra_pool_names,
storage_queue_sizes[4:],
)
}
self._drain_storage_control_queues_impl(
n_revoke=n_revoke,
n_storage_hit=n_storage_hit,
n_backup=n_backup,
n_release=n_release,
extra_release_counts=extra_release_counts,
log_metrics=True,
)
if self.enable_storage_metrics and self.storage_metrics_collector is not None:
self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats()
)
def flush_write_through_acks(self) -> None:
"""Flush pending write-through acknowledgements."""
self.writing_check()
def ready_to_load_host_cache(self) -> int:
"""Notify the cache controller to start the KV cache loading."""
if self.cache_controller is not None:
python/sglang/srt/session/streaming_session.py
-3
View File
@@ -602,9 +602,6 @@ class StreamingSession(BasePrefixCache):
def ready_to_load_host_cache(self):
return self.inner.ready_to_load_host_cache()
def flush_write_through_acks(self) -> None:
return self.inner.flush_write_through_acks()
def check_hicache_events(self):
return self.inner.check_hicache_events()