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

[HiCache] Keep hybrid transfer layer maps stage-local under PP (#39699)

Browse Source 
Co-authored-by: Aurick Qiao <6137920+aurickq@users.noreply.github.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
This commit is contained in:
Aurick Qiao
2026-09-17 10:18:44 +08:00
committed by GitHub
co-authored by Aurick Qiao Ke Bao
parent 0daa040e48
commit 4793f56835
2 changed files with 109 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+26 -5
View File
@@ -76,6 +76,15 @@ def _make_layer_mapper(
return mapper
def _stage_local_layer_mapping(
layer_mapping: dict[int, int], start_layer: int
) -> dict[int, int]:
return {
global_layer - start_layer: pool_layer
for global_layer, pool_layer in layer_mapping.items()
}
def _with_mtp_layer_mapping(
layer_mapping: dict[int, int],
*,
@@ -1478,8 +1487,12 @@ class _MambaStrategy(StackStrategy):
model_name=None,
enable_storage_metrics=False,
):
full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping)
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
full_layer_mapping = _stage_local_layer_mapping(
kvcache.full_attention_layer_id_mapping, kvcache.start_layer
)
mamba_layer_mapping = _stage_local_layer_mapping(
params.req_to_token_pool.mamba_map, kvcache.start_layer
)
host_pool_group, cache_controller = build_hybrid_mamba_stack(
params=params,
kv_pool=kvcache.full_kv_pool,
@@ -1511,9 +1524,15 @@ class _MambaStrategy(StackStrategy):
def _swa_layer_mappings(kvcache) -> tuple[dict[int, int], dict[int, int]]:
full = {
gid: lid for gid, (lid, is_swa) in kvcache.layers_mapping.items() if not is_swa
gid - kvcache.start_layer: lid
for gid, (lid, is_swa) in kvcache.layers_mapping.items()
if not is_swa
}
swa = {
gid - kvcache.start_layer: lid
for gid, (lid, is_swa) in kvcache.layers_mapping.items()
if is_swa
}
swa = {gid: lid for gid, (lid, is_swa) in kvcache.layers_mapping.items() if is_swa}
return full, swa
@@ -1600,7 +1619,9 @@ class _MambaSwaStrategy(StackStrategy):
enable_storage_metrics=False,
):
full_layer_mapping, swa_layer_mapping = _swa_layer_mappings(kvcache)
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
mamba_layer_mapping = _stage_local_layer_mapping(
params.req_to_token_pool.mamba_map, kvcache.start_layer
)
host_pool_group, cache_controller = build_hybrid_mamba_swa_stack(
params=params,
full_kv_pool=kvcache.full_kv_pool,