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 7fd9580b2..1319031f6 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 @@ -162,6 +162,8 @@ class HybridCacheController(BaseHiCacheController): 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, transfer_layer_num: Optional[int] = None, enable_storage_metrics: bool = False, ): @@ -180,6 +182,8 @@ 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 99198dfec..27d2b60d9 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 @@ -92,6 +92,8 @@ def build_nsa_hybrid_stack( storage_backend_extra_config=extra_config, pp_rank=radix_cache.pp_rank, pp_size=radix_cache.pp_size, + attn_cp_rank=params.attn_cp_rank, + attn_cp_size=params.attn_cp_size, transfer_layer_num=layer_num, enable_storage_metrics=enable_storage_metrics, ) @@ -190,6 +192,8 @@ def build_mamba_hybrid_stack( storage_backend_extra_config=extra_config, pp_rank=params.pp_rank, pp_size=params.pp_size, + attn_cp_rank=params.attn_cp_rank, + attn_cp_size=params.attn_cp_size, transfer_layer_num=transfer_layer_num, enable_storage_metrics=enable_storage_metrics, )