[HiCache]Fix CP support for hybrid model (#22782)
Co-authored-by: hzh0425 <hzh0425@apache.org>
This commit is contained in:
@@ -162,6 +162,8 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
storage_backend_extra_config: Optional[dict] = None,
|
storage_backend_extra_config: Optional[dict] = None,
|
||||||
pp_rank: int = 0,
|
pp_rank: int = 0,
|
||||||
pp_size: int = 1,
|
pp_size: int = 1,
|
||||||
|
attn_cp_rank: int = 0,
|
||||||
|
attn_cp_size: int = 1,
|
||||||
transfer_layer_num: Optional[int] = None,
|
transfer_layer_num: Optional[int] = None,
|
||||||
enable_storage_metrics: bool = False,
|
enable_storage_metrics: bool = False,
|
||||||
):
|
):
|
||||||
@@ -180,6 +182,8 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
storage_backend_extra_config=storage_backend_extra_config,
|
storage_backend_extra_config=storage_backend_extra_config,
|
||||||
pp_rank=pp_rank,
|
pp_rank=pp_rank,
|
||||||
pp_size=pp_size,
|
pp_size=pp_size,
|
||||||
|
attn_cp_rank=attn_cp_rank,
|
||||||
|
attn_cp_size=attn_cp_size,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
)
|
)
|
||||||
# Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)),
|
# Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)),
|
||||||
|
|||||||
@@ -92,6 +92,8 @@ def build_nsa_hybrid_stack(
|
|||||||
storage_backend_extra_config=extra_config,
|
storage_backend_extra_config=extra_config,
|
||||||
pp_rank=radix_cache.pp_rank,
|
pp_rank=radix_cache.pp_rank,
|
||||||
pp_size=radix_cache.pp_size,
|
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,
|
transfer_layer_num=layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
)
|
)
|
||||||
@@ -190,6 +192,8 @@ def build_mamba_hybrid_stack(
|
|||||||
storage_backend_extra_config=extra_config,
|
storage_backend_extra_config=extra_config,
|
||||||
pp_rank=params.pp_rank,
|
pp_rank=params.pp_rank,
|
||||||
pp_size=params.pp_size,
|
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,
|
transfer_layer_num=transfer_layer_num,
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
enable_storage_metrics=enable_storage_metrics,
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user