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

[HiCache] Add synchronization for context parallelism (#20460)

Browse Source 
Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
This commit is contained in:
Vladislav Nosivskoy
2026-04-28 02:13:07 +08:00
committed by GitHub
parent f34222da1b
commit 28ee08c172
9 changed files with 190 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/cache_controller.py
+67 -49
View File
@@ -253,6 +253,8 @@ class HiCacheController:
page_size: int,
tp_group: torch.distributed.ProcessGroup,
load_cache_event: threading.Event,
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
write_policy: str = "write_through_selective",
io_backend: str = "",
storage_backend: Optional[str] = None,
@@ -261,11 +263,12 @@ class HiCacheController:
storage_backend_extra_config: Optional[dict] = None,
pp_rank: int = 0,
pp_size: int = 1,
attn_cp_rank: int = 0,
attn_cp_size: int = 1,
enable_storage_metrics: bool = False,
):
self.tp_group = tp_group
self.attn_cp_group = attn_cp_group
self.attn_tp_group = attn_tp_group
self.prefetch_sync_groups: List[torch.distributed.ProcessGroup] = []
self.mem_pool_device_allocator = token_to_kv_pool_allocator
mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
@@ -282,8 +285,6 @@ class HiCacheController:
self.storage_backend_type = None
self.pp_rank = pp_rank
self.pp_size = pp_size
self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = attn_cp_size
self.enable_storage_metrics = enable_storage_metrics
# Default storage page IO functions (may be overridden by attach).
@@ -337,6 +338,51 @@ class HiCacheController:
# Preserve the historical error shape on init for unknown backends.
raise ValueError(f"Failed to create storage backend: {e}") from e
def get_attn_cp_rank_and_size(self) -> tuple[int, int]:
"""Derive CP rank/size from the attn_cp process group."""
if self.attn_cp_group is not None:
return (
torch.distributed.get_rank(group=self.attn_cp_group),
torch.distributed.get_world_size(group=self.attn_cp_group),
)
return 0, 1
def _create_prefetch_sync_groups(self) -> None:
from sglang.srt.distributed.parallel_state import create_custom_parallel_group
self.prefetch_sync_groups = []
seen_rank_sets = set()
if self.attn_cp_group is not None or self.attn_tp_group is not None:
base_groups = [self.attn_cp_group, self.attn_tp_group]
else:
base_groups = [self.tp_group]
for group in base_groups:
if group is None or torch.distributed.get_world_size(group=group) == 1:
continue
group_ranks = tuple(torch.distributed.get_process_group_ranks(group))
if group_ranks in seen_rank_sets:
continue
seen_rank_sets.add(group_ranks)
self.prefetch_sync_groups.append(
create_custom_parallel_group(
group_ranks=list(group_ranks), backend="gloo"
)
)
def _destroy_prefetch_sync_groups(self) -> None:
for group in self.prefetch_sync_groups:
try:
torch.distributed.destroy_process_group(group)
except Exception:
pass
self.prefetch_sync_groups = []
def _all_reduce_prefetch_groups(self, tensor: torch.Tensor, op) -> None:
for group in self.prefetch_sync_groups:
torch.distributed.all_reduce(tensor, op=op, group=group)
def _start_storage_threads(self):
"""Start storage prefetch/backup threads and their queues.
@@ -467,17 +513,9 @@ class HiCacheController:
# tracking the number of tokens locked in prefetching, updated by the main scheduler thread
self.prefetch_tokens_occupied = 0
# create a new communication group for synchronizing storage operations across TP workers
self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
if self.tp_world_size > 1:
from sglang.srt.distributed.parallel_state import (
create_custom_parallel_group,
)
group_ranks = torch.distributed.get_process_group_ranks(self.tp_group)
self.prefetch_tp_group = create_custom_parallel_group(
group_ranks=group_ranks, backend="gloo"
)
# Use dedicated gloo groups so storage prefetch sync is isolated
# from other collectives and consistent across CPxTP participants.
self._create_prefetch_sync_groups()
# Select the get and set functions
self.page_get_func = self._generic_page_get
@@ -502,15 +540,7 @@ class HiCacheController:
self._stop_storage_threads()
except Exception:
pass
try:
if hasattr(self, "prefetch_tp_group"):
try:
torch.distributed.destroy_process_group(self.prefetch_tp_group)
except Exception:
pass
self.prefetch_tp_group = None
except Exception:
pass
self._destroy_prefetch_sync_groups()
try:
if (
hasattr(self, "storage_backend")
@@ -547,19 +577,8 @@ class HiCacheController:
# to avoid flipping `enable_storage` flags while threads are still alive.
raise RuntimeError("Stop storage threads failed; detach aborted.") from e
# Best-effort destroy process group created for storage ops.
try:
if (
hasattr(self, "prefetch_tp_group")
and self.prefetch_tp_group is not None
):
try:
torch.distributed.destroy_process_group(self.prefetch_tp_group)
except Exception:
pass
self.prefetch_tp_group = None
except Exception:
pass
# Best-effort destroy process groups created for storage ops.
self._destroy_prefetch_sync_groups()
# Best-effort close (some backends rely on GC/destructor).
try:
@@ -613,13 +632,15 @@ class HiCacheController:
and tp_lcm_size > self.tp_size
)
attn_cp_rank, attn_cp_size = self.get_attn_cp_rank_and_size()
return HiCacheStorageConfig(
tp_rank=self.tp_rank,
tp_size=self.tp_size,
pp_rank=self.pp_rank,
pp_size=self.pp_size,
attn_cp_rank=self.attn_cp_rank,
attn_cp_size=self.attn_cp_size,
attn_cp_rank=attn_cp_rank,
attn_cp_size=attn_cp_size,
is_mla_model=is_mla_backend,
enable_storage_metrics=self.enable_storage_metrics,
is_page_first_layout=self.mem_pool_host.layout == "page_first",
@@ -963,16 +984,13 @@ class HiCacheController:
if operation is None:
continue
hash_value, storage_hit_count = self._storage_hit_query(operation)
if self.tp_world_size > 1:
storage_hit_count_tensor = torch.tensor(
storage_hit_count, dtype=torch.int
)
torch.distributed.all_reduce(
storage_hit_count_tensor,
op=torch.distributed.ReduceOp.MIN,
group=self.prefetch_tp_group,
)
storage_hit_count = storage_hit_count_tensor.item()
storage_hit_count_tensor = torch.tensor(
storage_hit_count, dtype=torch.int
)
self._all_reduce_prefetch_groups(
storage_hit_count_tensor, torch.distributed.ReduceOp.MIN
)
storage_hit_count = storage_hit_count_tensor.item()
if storage_hit_count < self.prefetch_threshold:
# not to prefetch if not enough benefits
python/sglang/srt/managers/scheduler.py
+2 -2
View File
@@ -803,14 +803,14 @@ class Scheduler(
if self.server_args.enable_dp_attention
else self.tp_cpu_group
),
attn_cp_cache_group=self.attn_cp_cpu_group,
attn_tp_cache_group=self.attn_tp_cpu_group,
eviction_policy=server_args.radix_eviction_policy,
enable_metrics=self.enable_metrics,
enable_kv_cache_events=self.enable_kv_cache_events,
enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
pp_rank=self.pp_rank,
pp_size=self.pp_size,
attn_cp_rank=self.attn_cp_rank,
attn_cp_size=self.attn_cp_size,
chunked_prefill_size=effective_chunked_prefill_size,
sliding_window_size=self.sliding_window_size,
)
python/sglang/srt/mem_cache/cache_init_params.py
+2
View File
@@ -20,6 +20,8 @@ class CacheInitParams:
is_eagle: bool = False
tp_cache_group: Optional[torch.distributed.ProcessGroup] = None
attn_cp_cache_group: Optional[torch.distributed.ProcessGroup] = None
attn_tp_cache_group: Optional[torch.distributed.ProcessGroup] = None
eviction_policy: str = "lru"
disable_finished_insert: bool = False
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+2
View File
@@ -143,6 +143,8 @@ class HiMambaRadixCache(MambaRadixCache):
prefetch_threshold=prefetch_threshold,
load_cache_event=self.load_cache_event,
enable_storage_metrics=self.enable_storage_metrics,
attn_cp_group=params.attn_cp_cache_group,
attn_tp_group=params.attn_tp_cache_group,
)
self._apply_storage_runtime_config(
storage_backend=server_args.hicache_storage_backend,
python/sglang/srt/mem_cache/hiradix_cache.py
+46 -42
View File
@@ -98,11 +98,11 @@ class HiRadixCache(RadixCache):
)
self.tp_group = params.tp_cache_group
self.attn_cp_group = params.attn_cp_cache_group
self.attn_tp_group = params.attn_tp_cache_group
self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
self.pp_rank = params.pp_rank
self.pp_size = params.pp_size
self.attn_cp_rank = params.attn_cp_rank
self.attn_cp_size = params.attn_cp_size
self.enable_storage = server_args.hicache_storage_backend is not None
self.enable_storage_metrics = self.enable_storage and params.enable_metrics
self.extra_metric_labels = server_args.extra_metric_labels
@@ -130,6 +130,8 @@ class HiRadixCache(RadixCache):
prefetch_threshold=prefetch_threshold,
enable_storage_metrics=self.enable_storage_metrics,
load_cache_event=self.load_cache_event,
attn_cp_group=self.attn_cp_group,
attn_tp_group=self.attn_tp_group,
)
else:
self.cache_controller = HiCacheController(
@@ -138,6 +140,8 @@ class HiRadixCache(RadixCache):
self.page_size,
self.tp_group,
load_cache_event=self.load_cache_event,
attn_cp_group=self.attn_cp_group,
attn_tp_group=self.attn_tp_group,
write_policy=server_args.hicache_write_policy,
io_backend=server_args.hicache_io_backend,
storage_backend=server_args.hicache_storage_backend,
@@ -146,8 +150,6 @@ class HiRadixCache(RadixCache):
storage_backend_extra_config=extra_config,
pp_rank=self.pp_rank,
pp_size=self.pp_size,
attn_cp_rank=self.attn_cp_rank,
attn_cp_size=self.attn_cp_size,
enable_storage_metrics=self.enable_storage_metrics,
)
self._apply_storage_runtime_config(
@@ -184,6 +186,24 @@ class HiRadixCache(RadixCache):
super().__init__(params=params)
def _all_reduce_attn_groups(self, tensor: torch.Tensor, op):
reduced = False
for group in (self.attn_cp_group, self.attn_tp_group):
if group is not None and torch.distributed.get_world_size(group=group) > 1:
torch.distributed.all_reduce(tensor, op=op, group=group)
reduced = True
if not reduced and self.tp_world_size > 1:
torch.distributed.all_reduce(tensor, op=op, group=self.tp_group)
def _barrier_attn_groups(self):
waited = False
for group in (self.attn_cp_group, self.attn_tp_group):
if group is not None and torch.distributed.get_world_size(group=group) > 1:
torch.distributed.barrier(group=group)
waited = True
if not waited and self.tp_world_size > 1:
torch.distributed.barrier(group=self.tp_group)
def shutdown(self):
"""Best-effort auto-detach of storage backend on process shutdown.
@@ -220,14 +240,17 @@ class HiRadixCache(RadixCache):
self.enable_storage_metrics = enable_storage_metrics
if self.enable_storage_metrics:
attn_cp_rank, attn_cp_size = (
self.cache_controller.get_attn_cp_rank_and_size()
)
labels = {
"storage_backend": storage_backend,
"tp_rank": self.cache_controller.tp_rank,
"dp_rank": self.cache_controller.dp_rank,
"pp_rank": self.cache_controller.pp_rank,
"pp_size": self.cache_controller.pp_size,
"attn_cp_rank": self.cache_controller.attn_cp_rank,
"attn_cp_size": self.cache_controller.attn_cp_size,
"attn_cp_rank": attn_cp_rank,
"attn_cp_size": attn_cp_size,
}
if extra_metric_labels:
labels.update(extra_metric_labels)
@@ -741,13 +764,8 @@ class HiRadixCache(RadixCache):
break
finish_count += 1
queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
if self.tp_world_size > 1:
# synchronize TP workers to make the same update to radix cache
torch.distributed.all_reduce(
queue_size,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
)
# Keep cache state transitions identical across CPxTP participants.
self._all_reduce_attn_groups(queue_size, torch.distributed.ReduceOp.MIN)
finish_count = int(queue_size.item())
while finish_count > 0:
@@ -1074,10 +1092,7 @@ class HiRadixCache(RadixCache):
],
dtype=torch.int,
)
if self.tp_world_size > 1:
torch.distributed.all_reduce(
qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
)
self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN)
n_revoke, n_backup, n_release = map(int, qsizes.tolist())
self._drain_storage_control_queues_impl(
@@ -1118,18 +1133,13 @@ class HiRadixCache(RadixCache):
return True
operation_terminated = operation.is_terminated()
if self.tp_world_size > 1:
states = torch.tensor(
[1 - int(can_terminate), int(operation_terminated)],
dtype=torch.int,
)
torch.distributed.all_reduce(
states,
op=torch.distributed.ReduceOp.MAX,
group=self.tp_group,
)
can_terminate = states[0].item() == 0
operation_terminated = states[1].item() == 1
states = torch.tensor(
[1 - int(can_terminate), int(operation_terminated)],
dtype=torch.int,
)
self._all_reduce_attn_groups(states, torch.distributed.ReduceOp.MAX)
can_terminate = states[0].item() == 0
operation_terminated = states[1].item() == 1
# the operation should be terminated if it is already terminated on any TP worker
# or it meets the termination condition on all TP workers
can_terminate = can_terminate or operation_terminated
@@ -1159,17 +1169,12 @@ class HiRadixCache(RadixCache):
logger.debug(f"Prefetch {req_id} completed with {completed_tokens} tokens")
min_completed_tokens = completed_tokens
if self.tp_world_size > 1:
# synchrnoize TP workers to make the same update to hiradix cache
completed_tokens_tensor = torch.tensor(
min_completed_tokens, dtype=torch.int
)
torch.distributed.all_reduce(
completed_tokens_tensor,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
)
min_completed_tokens = completed_tokens_tensor.item()
# Synchronize workers before mutating host cache tree state.
completed_tokens_tensor = torch.tensor(min_completed_tokens, dtype=torch.int)
self._all_reduce_attn_groups(
completed_tokens_tensor, torch.distributed.ReduceOp.MIN
)
min_completed_tokens = completed_tokens_tensor.item()
fetched_token_ids = token_ids[:min_completed_tokens]
written_indices = host_indices[:min_completed_tokens]
matched_length = self._insert_helper_host(
@@ -1494,8 +1499,7 @@ class HiRadixCache(RadixCache):
return
completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
if self.tp_world_size > 1:
torch.distributed.barrier(group=self.tp_group)
self._barrier_attn_groups()
last_host_node.release_host()
del self.ongoing_prefetch[rid]
self.cache_controller.append_host_mem_release(host_indices[:completed_tokens])
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+4 -2
View File
@@ -154,6 +154,8 @@ class HybridCacheController(BaseHiCacheController):
page_size: int,
tp_group: torch.distributed.ProcessGroup,
load_cache_event: threading.Event,
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
write_policy: str = "write_through_selective",
io_backend: str = "",
storage_backend: Optional[str] = None,
@@ -174,6 +176,8 @@ class HybridCacheController(BaseHiCacheController):
page_size=page_size,
tp_group=tp_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
write_policy=write_policy,
io_backend=io_backend,
storage_backend=None,
@@ -182,8 +186,6 @@ class HybridCacheController(BaseHiCacheController):
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=pp_rank,
pp_size=pp_size,
attn_cp_rank=attn_cp_rank,
attn_cp_size=attn_cp_size,
enable_storage_metrics=enable_storage_metrics,
)
# Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)),
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+28
View File
@@ -17,6 +17,8 @@ from sglang.srt.mem_cache.memory_pool_host import (
)
if TYPE_CHECKING:
import torch
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
@@ -94,6 +96,8 @@ def build_kv_only_stack(
page_size: int,
tp_group,
load_cache_event,
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
storage_backend: Optional[str],
use_mla: bool,
override_kv_cache_dim: Optional[int] = None,
@@ -131,6 +135,8 @@ def build_kv_only_stack(
page_size,
tp_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
write_policy=server_args.hicache_write_policy,
io_backend=server_args.hicache_io_backend,
storage_backend=storage_backend,
@@ -158,6 +164,8 @@ def build_hybrid_mamba_stack(
page_size: int,
tp_group,
load_cache_event,
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
storage_backend: Optional[str],
use_mla: bool,
host_mamba_evict_fn: Optional[Callable[[int], Any]] = None,
@@ -211,6 +219,8 @@ def build_hybrid_mamba_stack(
page_size,
tp_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
write_policy=server_args.hicache_write_policy,
io_backend=server_args.hicache_io_backend,
storage_backend=storage_backend,
@@ -237,6 +247,8 @@ def build_shared_anchor_stack(
page_size: int,
tp_group,
load_cache_event,
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
storage_backend: Optional[str],
use_mla: bool,
override_kv_cache_dim: Optional[int] = None,
@@ -284,6 +296,8 @@ def build_shared_anchor_stack(
page_size,
tp_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
write_policy=server_args.hicache_write_policy,
io_backend=server_args.hicache_io_backend,
storage_backend=storage_backend,
@@ -306,6 +320,8 @@ def attach_hybrid_pool_to_unified_cache(
server_args: ServerArgs,
*,
load_cache_event,
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
) -> None:
"""Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache."""
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
@@ -342,6 +358,8 @@ def attach_hybrid_pool_to_unified_cache(
page_size=cache.page_size,
tp_group=params.tp_cache_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
use_mla=use_mla,
host_mamba_evict_fn=lambda n: cache.evict_host(n, ComponentType.MAMBA),
@@ -375,6 +393,8 @@ def attach_hybrid_pool_to_unified_cache(
page_size=cache.page_size,
tp_group=params.tp_cache_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=None,
use_mla=use_mla,
pp_rank=params.pp_rank,
@@ -411,6 +431,8 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
prefetch_threshold: int,
enable_storage_metrics: bool,
load_cache_event,
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
) -> None:
"""Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache.
@@ -428,6 +450,8 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
page_size=radix_cache.page_size,
tp_group=radix_cache.tp_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=server_args.hicache_storage_backend,
use_mla=True,
prefetch_threshold=prefetch_threshold,
@@ -467,6 +491,8 @@ def attach_hybrid_pool_to_mamba_cache(
prefetch_threshold: int,
load_cache_event,
enable_storage_metrics: bool = False,
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
) -> None:
"""Attach HostPoolGroup (KV + Mamba) + HybridCacheController for HiMambaRadixCache.
@@ -487,6 +513,8 @@ def attach_hybrid_pool_to_mamba_cache(
page_size=params.page_size,
tp_group=params.tp_cache_group,
load_cache_event=load_cache_event,
attn_cp_group=attn_cp_group,
attn_tp_group=attn_tp_group,
storage_backend=server_args.hicache_storage_backend,
use_mla=hybrid_kv.use_mla,
host_mamba_evict_fn=mamba_cache.evict_mamba_host,
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
+5 -9
View File
@@ -410,12 +410,9 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
self.attn_cp_size = 1
self.enable_pp = self.pp_size > 1
self.enable_cp = self.attn_cp_size > 1
if self.enable_pp or self.enable_cp:
self.mha_suffix = (
f"{self.local_rank}_{self.pp_rank}_{self.attn_cp_rank}"
)
self.mla_suffix = f"{self.pp_rank}_{self.attn_cp_rank}"
if self.enable_pp:
self.mha_suffix = f"{self.local_rank}_{self.pp_rank}"
self.mla_suffix = f"{self.pp_rank}"
else:
self.mha_suffix = f"{self.local_rank}"
self.mla_suffix = ""
@@ -428,10 +425,9 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
)
base_rank = self.local_rank * self.split_factor
target_ranks = [base_rank + i for i in range(self.split_factor)]
if self.enable_pp or self.enable_cp:
if self.enable_pp:
self.mha_suffix = [
f"{rank}_{self.pp_rank}_{self.attn_cp_rank}"
for rank in target_ranks
f"{rank}_{self.pp_rank}" for rank in target_ranks
]
else:
self.mha_suffix = [f"{rank}" for rank in target_ranks]