执行摘要
- 一句话：修复 DeepSeek V4 PP HiCache SWA 分配与层映射
- 推荐动作：面向阅读者的建议：该 PR 值得精读，尤其是关注流水线并行与缓存模块集成的开发者。其设计关键点在于：通过 PP 阶段本地层数分配资源，并通过间接方法封装索引转换，避免了全局层映射带来的状态不一致。讨论中关于近似算法精度的排查思路也值得借鉴。
功能与动机
DeepSeek V4 + HiCache 在流水线并行下会失败，因为 SWA KV pool 按全局模型层数分配，而每个 PP 阶段只拥有局部层切片。在 PP + HiCache 运行时，这会导致后续 PP 阶段的 HiCache 状态构建无效，并在服务初始化时失败，报错 ValueError: deepseek_v4_c4_state state_pools must not contain None。
实现拆解
- SWA KV 池按 PP 阶段层数分配：在
deepseek_v4_memory_pool.py 的 __init__ 中，计算 stage_layer_num = len(stage_ratios)，替换原来的全局 layer_num，作为 SWA KV 池的 layer_num 参数。
- HiCache 堆栈构建使用局部层映射：在
hybrid_pool_assembler.py 的 build_deepseek_v4_hicache_stack 中，移除 TODO 注释，添加 SWA 池缓冲区数量与传输层数的校验；遍历层映射时改用 local_layer_id 和 global_layer_id 分别构建 C4 状态本地列表和全局列表，SWA 层映射也基于局部数量构建。
- 新增局部索引方法：在
deepseek_v4_memory_pool.py 中新增 get_swa_raw_buffer(layer_id) 方法，通过 _swa_local_layer_id 将全局层 ID 转换为本地索引，返回对应的原始缓冲区。
- 模型前向路径适配：在
deepseek_v4.py 的两个前向准备函数中，将直接访问 swa_kv_pool.kv_buffer[self.layer_id] 替换为调用 get_swa_raw_buffer(self.layer_id)。
- 测试重构与新增：在基础测试类
test_unified_radix_cache_kl_dsv4.py 中抽取 _server_args 方法，支持通过类属性 pp_size/tp_size 动态设置服务器参数；新增 test_unified_radix_cache_kl_dsv4_pp.py 定义 TestUnifiedDeepSeekV4FlashHiCachePP4TP2 类，继承基础测试并覆盖 pp_size=4, tp_size=2，注册到 CI 的 extra-b 阶段（8-gpu-h200）。
关键文件：
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py（模块 缓存层；类别 source；类型 core-logic；符号 build_deepseek_v4_hicache_stack）: HiCache 堆栈构建的核心函数，修正了 SWA 和 C4 状态层映射从全局索引改为局部索引，并添加了池大小校验。python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py（模块 缓存层；类别 source；类型 core-logic；符号 get_swa_raw_buffer）: SWA KV 池分配参数改为使用 PP 阶段层数；新增 get_swa_raw_buffer 方法封装局部索引，是修复的关键数据结构变更。python/sglang/srt/models/deepseek_v4.py（模块 模型；类别 source；类型 data-contract；符号 _forward_prepare, _forward_prepare_multi_stream_hip）: 模型前向路径中直接访问 kv_buffer 的位置替换为 get_swa_raw_buffer，确保 PP 下索引正确。test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py（模块 测试；类别 test；类型 test-coverage；符号 _server_args）: 重构基础测试类，提取 _server_args 方法以支持 pp_size/tp_size 参数化，便于被 PP 测试类继承。test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4_pp.py（模块 测试；类别 test；类型 test-coverage；符号 TestUnifiedDeepSeekV4FlashHiCachePP4TP2）: 新增 PP4TP2 集成测试，验证 PP+HiCache 组合的正确性，填补 CI 覆盖空白。
关键符号：build_deepseek_v4_hicache_stack, get_swa_raw_buffer, _forward_prepare, _forward_prepare_multi_stream_hip, _server_args, test_multiturn_logprobs_match
关键源码片段
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
HiCache 堆栈构建的核心函数，修正了 SWA 和 C4 状态层映射从全局索引改为局部索引，并添加了池大小校验。
def build_deepseek_v4_hicache_stack(
...,
kvcache: Any,
page_size: int,
...
) -> tuple[HostPoolGroup, HybridCacheController]:
transfer_layer_num = kvcache.end_layer - kvcache.start_layer
full_layer_mapping = {layer_id: layer_id for layer_id in range(transfer_layer_num)}
# 验证 SWA KV 池的缓冲区数量必须等于本 PP 阶段的层数
if len(kvcache.swa_kv_pool.kv_buffer) != transfer_layer_num:
raise ValueError(
"DeepSeek V4 SWA KV pool must be PP-stage-local: "
f"got {len(kvcache.swa_kv_pool.kv_buffer)} buffers for "
f"{transfer_layer_num} local layers"
)
swa_layer_mapping = {layer_id: layer_id for layer_id in range(transfer_layer_num)}
c4_layer_mapping = {}
c128_layer_mapping = {}
c4_state_local_layers = []
c4_state_global_layers = []
for local_layer_id, layer_item in enumerate(
kvcache.layer_mapping[kvcache.start_layer : kvcache.end_layer]
):
global_layer_id = kvcache.start_layer + local_layer_id
if layer_item.compress_ratio == 4:
c4_layer_mapping[local_layer_id] = layer_item.compress_layer_id
c4_state_local_layers.append(local_layer_id)
c4_state_global_layers.append(global_layer_id)
elif layer_item.compress_ratio == 128:
c128_layer_mapping[local_layer_id] = layer_item.compress_layer_id
c4_state_mapping = {
layer_id: local_id for local_id, layer_id in enumerate(c4_state_local_layers)
}
# ... 后续构建 host pool 和 entries
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
SWA KV 池分配参数改为使用 PP 阶段层数；新增 get_swa_raw_buffer 方法封装局部索引，是修复的关键数据结构变更。
def __init__(self, ...):
# ...
stage_layer_num = len(stage_ratios) # PP 阶段的局部层数
c4_layer_num = sum(1 for r in stage_ratios if r == 4)
c128_layer_num = sum(1 for r in stage_ratios if r == 128)
# ...
# 非 unified_kv 路径：SWA KV 池使用 stage_layer_num 而非全局 layer_num
self.swa_kv_pool = self._make_kv_pool(
size=swa_size,
page_size=swa_page_size,
dtype=dtype,
layer_num=stage_layer_num, # 修正点
device=device,
enable_memory_saver=enable_memory_saver,
global_page_size=swa_page_size,
)
# ...
def get_swa_raw_buffer(self, layer_id: int) -> torch.Tensor:
"""根据全局层 ID 返回本阶段 SWA 池的原始缓冲区。"""
return self.swa_kv_pool.kv_buffer[self._swa_local_layer_id(layer_id)]
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py
重构基础测试类，提取 _server_args 方法以支持 pp_size/tp_size 参数化，便于被 PP 测试类继承。
@classmethod
def _server_args(cls):
args = [
"--trust-remote-code",
"--tp-size",
str(cls.tp_size),
]
if cls.pp_size != 1: # 当 PP 开启时添加 --pp-size 参数
args += ["--pp-size", str(cls.pp_size)]
args += [
# ... 其他固定参数
"--attention-backend",
"compressed",
"--page-size",
"256",
# ...
]
return args
评论区精华
风险与影响
- 风险：
- SWA 池分配变更：
stage_layer_num 在非 PP 时等于全局层数，行为兼容。但有潜在风险：若其他代码路径直接引用 swa_kv_pool.kv_buffer 而不使用 get_swa_raw_buffer，可能在 PP 下越界。需确保所有 SWA 访问都经过转换。
- 精度风险：KL 测试通过，但作者指出的 top-k 近似问题可能在其他场景下引发精度退化，建议用户关注环境变量后备机制。
- CI 覆盖：仅测试 PP4TP2 一种配置，未覆盖其他并行度组合，但基础测试已涵盖 TP4 非 PP 场景。
- 内存占用：SWA 池分配大小减小，不会出现内存浪费。
- 影响：
- 用户：启用 PP 和 HiCache 的 DeepSeek V4 场景不再崩溃，可以在流水线并行下正常使用分层缓存。
- 系统：SWA 池的正确分配修复了 HiCache 状态构建逻辑，PP 环境下的服务初始化顺利通过。
- 团队：新增的集成测试填补了 PP+HiCache 的 CI 空白，降低后续开发中的回归风险。
- 风险标记：SWA 池分配变更影响 PP/非 PP 路径, 近似 top-k 算法可能导致精度偏差, CI 测试覆盖仅限 PP4TP2 一种配置
关联脉络
- PR #28614 [HiCache] remove large host mem constraint: 同为 HiCache 相关修改，涉及内存池约束调整，与本 PR 共同完善 HiCache 功能。
参与讨论