diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 9a4160791..652b71b9a 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a7e902dcd..50797f988 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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, ) diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index 6f6fafae0..713bf308f 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -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 diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 455b8623d..3c2988854 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -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, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 6b2c1fe21..718d186c3 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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]) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 28a1062a5..86c5db24d 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -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)), diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 0bd0f037a..2b1ec8aac 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -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, diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 5ee3d24a5..9d65a7265 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -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] diff --git a/test/registered/hicache/test_hicache_storage_mooncake_backend.py b/test/registered/hicache/test_hicache_storage_mooncake_backend.py index 1a9b75a39..15d84a240 100644 --- a/test/registered/hicache/test_hicache_storage_mooncake_backend.py +++ b/test/registered/hicache/test_hicache_storage_mooncake_backend.py @@ -17,6 +17,7 @@ from sglang.test.test_utils import ( DEFAULT_MLA_MODEL_NAME_FOR_TEST, CustomTestCase, find_available_port, + get_gpu_count, is_in_ci, ) @@ -260,6 +261,39 @@ class TestMooncakeBackendMLAModel( return server_args, env_vars +@unittest.skipUnless(get_gpu_count() >= 2, "Requires at least 2 CUDA GPUs for TP2+CP2") +class TestMooncakeBackendQwen330BCP2( + HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase +): + """Qwen3-30B with Mooncake HiCache storage, CP2, and TP2.""" + + @classmethod + def _get_model_name(cls): + return "Qwen/Qwen3-30B-A3B-FP8" + + @classmethod + def _get_additional_server_args_and_env(cls): + server_args, env_vars = super()._get_additional_server_args_and_env() + server_args.update( + { + "--tp-size": 2, + "--moe-dp-size": 2, + "--attn-cp-size": 2, + "--enable-prefill-context-parallel": True, + "--trust-remote-code": True, + "--cuda-graph-max-bs": 32, + "--max-running-requests": 32, + "--max-total-tokens": 8192, + "--model-loader-extra-config": ( + '{"enable_multithread_load": true, "num_threads": 64}' + ), + "--hicache-mem-layout": "page_first_direct", + "--hicache-io-backend": "direct", + } + ) + return server_args, env_vars + + class TestMooncakeBackendAccuracy( HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase ):