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

[UnifiedTree]: Fix CP Reduce For L3 HiCache (#27330)

Browse Source
This commit is contained in:
Zhangheng
2026-06-05 14:03:54 +08:00
committed by GitHub
parent bd47869ba4
commit 4df1ccdadc
1 changed files with 4 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_radix_cache.py
+4 -15
View File
@@ -1861,10 +1861,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
[1 - int(can_terminate), int(operation_terminated)], [1 - int(can_terminate), int(operation_terminated)],
dtype=torch.int, dtype=torch.int,
) )
if self.tp_world_size > 1: self._all_reduce_attn_groups(states, torch.distributed.ReduceOp.MAX)
torch.distributed.all_reduce(
states, op=torch.distributed.ReduceOp.MAX, group=self.tp_group
)
can_terminate = states[0].item() == 0 can_terminate = states[0].item() == 0
operation_terminated = states[1].item() == 1 operation_terminated = states[1].item() == 1
return can_terminate or operation_terminated return can_terminate or operation_terminated
@@ -1899,11 +1896,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
[completed_tokens] + [hit_pages.get(p, 0) for p in sidecar_pools], [completed_tokens] + [hit_pages.get(p, 0) for p in sidecar_pools],
dtype=torch.int, dtype=torch.int,
) )
torch.distributed.all_reduce( self._all_reduce_attn_groups(packed, torch.distributed.ReduceOp.MIN)
packed,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
)
min_completed_tokens = int(packed[0].item()) min_completed_tokens = int(packed[0].item())
for i, p in enumerate(sidecar_pools, start=1): for i, p in enumerate(sidecar_pools, start=1):
hit_pages[p] = int(packed[i].item()) hit_pages[p] = int(packed[i].item())
@@ -1979,8 +1972,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
return return
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
if self.tp_world_size > 1: self._barrier_attn_groups()
torch.distributed.barrier(group=self.tp_group)
self.dec_host_lock_ref(last_host_node, anchor_lock_params) self.dec_host_lock_ref(last_host_node, anchor_lock_params)
del self.ongoing_prefetch[rid] del self.ongoing_prefetch[rid]
self.cache_controller.append_host_mem_release( self.cache_controller.append_host_mem_release(
@@ -2103,10 +2095,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
local_qsize_list, local_qsize_list,
dtype=torch.int, dtype=torch.int,
) )
if self.tp_world_size > 1: self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN)
torch.distributed.all_reduce(
qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
)
qsize_list = list(map(int, qsizes.tolist())) qsize_list = list(map(int, qsizes.tolist()))
n_revoke, n_backup, n_release = qsize_list[:3] n_revoke, n_backup, n_release = qsize_list[:3]
extra_release_counts = { extra_release_counts = {