已开启

HP DeepSeek V4 CP 后端适配组件化设计 #3

yide12创建于 6月11日

6月11日创建

HP DeepSeek V4 CP 后端适配组件化设计

本文整理当前 DeepSeek V4 CP local-tensor 实现的下一步组件化方向。核心目标是：

HP 不只支持 MindFormers 后端，也要支持 Torch 后端。

因此 DeepSeek V4 CP 不能继续以 MF 的 parent module 为中心设计，而应该拆成：

后端无关 CP 算法层
+
后端/模型适配层

当前代码位置：

/home/wyd/code_review_deepseek_v4_cp/hp_feature/hyper_parallel/core/context_parallel/deepseek_v4_context_parallel.py

1. 当前实现的问题

当前实现逻辑正确，但 CP forward 强依赖 MF CompressedSparseAttention module：

self.parent.compressor
self.parent.indexer.compressor
self.parent._construct_naive
self.parent.squeeze
self.parent.window_size
self.parent.compress_ratio
self.parent.force_unfused

这对 MF 很自然，因为这些属性都在 MF CSA module 上。

但对 Torch 后端来说，这些名字、forward 签名、attention 入口、compressor 封装方式可能完全不同。如果 CP 核心逻辑继续直接访问 parent.xxx，后续 Torch 支持会变成大量 if backend == "torch" 分支。

所以需要把 “CP 算法” 和 “后端模型对象怎么调用” 分离。

2. 推荐总体结构

flowchart TD
    MF[MindFormers CSA Module] --> MFA[MindFormers Adapter]
    TH[Torch CSA Module / Functions] --> TA[Torch Adapter]

    MFA --> Core[DeepSeekV4 CP Core]
    TA --> Core

    Core --> C1[CP Meta]
    Core --> C2[Window KV Halo]
    Core --> C3[Ratio4 X Halo]
    Core --> C4[Global Compressed KV]
    Core --> C5[Select And Pad]
    Core --> C6[Backend Attention Call]

    C6 --> MFA
    C6 --> TA

设计原则：

CP Core 只管算法阶段。
Adapter 只管后端模块怎么取属性、怎么调用 compressor、怎么调用 attention。

3. 两层职责边界

3.1 CP Core 负责

CP Core 后端无关，负责：

根据 cp_mesh 和 query 推导 CP meta。
准备 x boundary halo。
准备 window ori_kv left halo。
调用 adapter 提供的 compressor 生成 local compressed slots。
all-gather compressed slots。
select-and-pad 得到 global compressed KV。
将准备好的 ori_ext / cmp_kv_global / cmp_idx_global / pos_offset / ori_halo_len 交给 adapter 的 attention 入口。

CP Core 不应该直接知道：

parent.compressor
parent.indexer
parent._construct_naive
parent.squeeze
MindFormers Cell
Torch Module

3.2 Adapter 负责

Adapter 负责屏蔽后端差异。

它需要告诉 CP Core：

当前 layer 的 compress_ratio。
当前 layer 的 window_size。
是否允许 small-op/naive 路径。
怎么从 key 得到 ori_kv。
怎么调用主 compressor。
怎么调用 ratio=4 indexer compressor。
怎么调用后端 attention。

4. Adapter 协议

建议定义一个轻量协议，不一定一开始就用 abc.ABC，但语义应稳定。

class DeepSeekV4CSAAdapter:
    @property
    def compress_ratio(self) -> int:
        ...

    @property
    def window_size(self) -> int:
        ...

    @property
    def force_unfused(self) -> bool:
        ...

    @property
    def has_indexer(self) -> bool:
        ...

    def extract_ori_kv(self, key):
        ...

    def compress_main(self, x, halo_x, cp_rank: int, *, overlap: bool):
        ...

    def compress_indexer(self, x, halo_x, cp_rank: int):
        ...

    def run_attention(
            self,
            query,
            ori_ext,
            x,
            qr,
            *,
            pos_offset: int,
            cmp_kv_global,
            ori_halo_len: int,
            cmp_idx_global,
    ):
        ...

这个协议的关键点是：

CP Core 不再持有 MF parent。
CP Core 只持有 adapter。

5. MindFormers Adapter

MF adapter 负责把 MF CompressedSparseAttention module 包成统一接口。

示意：

class MindFormersDeepSeekV4CSAAdapter:
    def __init__(self, csa_module):
        self.module = csa_module

    @property
    def compress_ratio(self):
        return int(getattr(self.module, "compress_ratio", 0) or 0)

    @property
    def window_size(self):
        return int(getattr(self.module, "window_size", 128))

    @property
    def force_unfused(self):
        return bool(getattr(self.module, "force_unfused", True))

    @property
    def has_indexer(self):
        return getattr(self.module, "indexer", None) is not None

    def extract_ori_kv(self, key):
        if len(key.shape) == 3:
            return key
        return self.module.squeeze(key, 2)

    def compress_main(self, x, halo_x, cp_rank, *, overlap):
        return stage1_compress_with_halo(
            self.module.compressor,
            x,
            halo_x,
            cp_rank,
            overlap=overlap,
        )

    def compress_indexer(self, x, halo_x, cp_rank):
        return stage1_compress_with_halo(
            self.module.indexer.compressor,
            x,
            halo_x,
            cp_rank,
            overlap=True,
        )

    def run_attention(self, query, ori_ext, x, qr, **kwargs):
        return self.module._construct_naive(
            query,
            ori_ext,
            x,
            qr,
            **kwargs,
        )

MF patch API 可以保留：

def patch_deepseek_v4_compressed_sparse_attention_forward(
        csa_module,
        *,
        cp_mesh,
        config=None,
):
    adapter = MindFormersDeepSeekV4CSAAdapter(csa_module)
    cp_forward = DeepSeekV4CPForward(adapter, cp_mesh, config)
    ...

这样对 MF 侧接口不变，但内部不再让 CP Core 直接读 parent。

6. Torch Adapter

Torch adapter 可以包 torch module，也可以包一组函数。

示意：

class TorchDeepSeekV4CSAAdapter:
    def __init__(
            self,
            module,
            *,
            compressor,
            indexer_compressor=None,
            attention_fn,
            compress_ratio: int,
            window_size: int,
            force_unfused: bool = True,
    ):
        self.module = module
        self.compressor = compressor
        self.indexer_compressor = indexer_compressor
        self.attention_fn = attention_fn
        self._compress_ratio = compress_ratio
        self._window_size = window_size
        self._force_unfused = force_unfused

    @property
    def compress_ratio(self):
        return self._compress_ratio

    @property
    def window_size(self):
        return self._window_size

    @property
    def force_unfused(self):
        return self._force_unfused

    @property
    def has_indexer(self):
        return self.indexer_compressor is not None

    def extract_ori_kv(self, key):
        return key if key.ndim == 3 else key.squeeze(2)

    def compress_main(self, x, halo_x, cp_rank, *, overlap):
        return stage1_compress_with_halo(
            self.compressor,
            x,
            halo_x,
            cp_rank,
            overlap=overlap,
        )

    def compress_indexer(self, x, halo_x, cp_rank):
        return stage1_compress_with_halo(
            self.indexer_compressor,
            x,
            halo_x,
            cp_rank,
            overlap=True,
        )

    def run_attention(self, query, ori_ext, x, qr, **kwargs):
        return self.attention_fn(
            query=query,
            ori_kv=ori_ext,
            x=x,
            qr=qr,
            **kwargs,
        )

Torch 后端的关键是：

不要要求 Torch module 长得像 MF module。
只要求 Torch adapter 实现相同协议。

7. CP Core 主流程

重构后的 CP Core 主流程应类似：

class DeepSeekV4CPForward:
    def __init__(self, adapter, cp_mesh, config=None):
        self.adapter = adapter
        self.cp_mesh = cp_mesh
        self.config = config or DeepSeekV4CPConfig()

    def __call__(self, query, key, value, attention_mask, x=None, qr=None, **kwargs):
        if self.cp_mesh.size() <= 1:
            return self.adapter.run_origin(
                query, key, value, attention_mask, x=x, qr=qr, **kwargs
            )

        cp_ctx = self._build_cp_context(query)
        self._validate_inputs(query, key, x, qr, cp_ctx)
        return self._construct_cp(query, key, x, qr, cp_ctx)

    def _construct_cp(self, query, key, x, qr, cp_ctx):
        ori_kv = self.adapter.extract_ori_kv(key)

        halo_x = build_ratio4_x_halo(
            x=x,
            cp_ctx=cp_ctx,
        )

        ori_ext, ori_halo_len = build_window_kv_for_cp(
            ori_kv=ori_kv,
            cp_ctx=cp_ctx,
        )

        cmp_kv_global = build_global_compressed_kv_for_cp(
            compress_fn=self.adapter.compress_main,
            x=x,
            halo_x=halo_x,
            cp_ctx=cp_ctx,
        )

        cmp_idx_global = None
        if cp_ctx.compress_ratio == 4 and self.adapter.has_indexer:
            cmp_idx_global = build_global_compressed_kv_for_cp(
                compress_fn=self.adapter.compress_indexer,
                x=x,
                halo_x=halo_x,
                cp_ctx=cp_ctx,
            )

        return self.adapter.run_attention(
            query,
            ori_ext,
            x,
            qr,
            pos_offset=cp_ctx.local_seq_start,
            cmp_kv_global=cmp_kv_global,
            ori_halo_len=ori_halo_len,
            cmp_idx_global=cmp_idx_global,
        )

主流程读起来应该是业务语义：

extract ori_kv
prepare x halo
prepare window KV halo
prepare global compressed KV
prepare global indexer KV
run attention

而不是暴露一堆通信细节。

8. 建议文件结构

可以先维持一个文件，稳定后再拆。最终建议结构：

hyper_parallel/core/context_parallel/deepseek_v4/
    __init__.py
    config.py
    forward.py
    collectives.py
    indices.py
    compression.py
    adapters/
        __init__.py
        mindformers.py
        torch.py

各文件职责：

config.py
    DeepSeekV4CPConfig
    DeepSeekV4CPContext

forward.py
    DeepSeekV4CPForward
    后端无关 CP 总控

collectives.py
    _all_gather_seq
    gather_x_boundary_halo
    gather_window_kv_tail
    gather_compressed_slots

indices.py
    build_select_pad_gather_index
    build_cmp_cu_seqlens
    select_and_pad_smallop

compression.py
    stage1_compress_with_halo
    build_global_compressed_kv_for_cp

adapters/mindformers.py
    MindFormersDeepSeekV4CSAAdapter
    patch_mindformers_deepseek_v4_csa_forward

adapters/torch.py
    TorchDeepSeekV4CSAAdapter
    patch_torch_deepseek_v4_csa_forward

9. Public API 建议

短期兼容 MF：

patch_deepseek_v4_compressed_sparse_attention_forward(
    csa_module,
    *,
    cp_mesh,
    config=None,
)

内部改成 MF adapter。

中期可以显式区分后端：

patch_mindformers_deepseek_v4_csa_forward(
    csa_module,
    *,
    cp_mesh,
    config=None,
)

patch_torch_deepseek_v4_csa_forward(
    module,
    *,
    cp_mesh,
    config=None,
    compressor=None,
    indexer_compressor=None,
    attention_fn=None,
)

不建议一开始做：

patch_deepseek_v4_csa_forward(..., backend="mindformers")

原因是 MF 和 Torch 的 module 形态、forward/construct 命名、训练 loss 处理不同。显式 API 比一个大函数里写很多 backend 分支更清楚。

10. Local Tensor 与 DTensor 的关系

当前设计仍然是 local tensor 版本。

不建议为了组件化强行 DTensor 化，原因是 DeepSeek V4 CP 不是简单均匀切分：

ori_ext = [left_halo, local_ori_kv] 不是普通 shard。
x_ext = [left_halo_x, local_x] 不是普通 shard。
ratio=4 的 boundary helper slot 需要 all-gather 后 select-and-pad。
packed/TND 还需要样本边界和 compressed cu_seqlens。

因此当前更合理的定位是：

后端无关 local-tensor CP algorithm
+
显式通信 collectives
+
后端 adapter

DTensor/DistributedOp 可以作为后续 fused op 路径的能力补充，不是首版 CP Core 的基础表达。

11. 迁移步骤建议

Step 1：先在当前文件内抽象 adapter

不急着拆文件，先把 self.parent.xxx 收敛到 adapter。

目标：

DeepSeekV4CompressedSparseAttentionCPForward 不再直接访问 parent。

Step 2：把主流程改成语义 helper

将 _construct_cp 改成：

prepare_ratio4_x_halo
prepare_window_kv
prepare_global_compressed_kv
prepare_global_indexer_kv
run_attention

Step 3：补 Torch adapter

先支持最小 Torch fake module / fake attention_fn，跑通 shape demo。

Step 4：拆目录

当 MF adapter 和 Torch adapter 都能跑通后，再拆成 deepseek_v4/ 子目录。

Step 5：扩展 fused 路径

后续再把 PR717 fused shared-KV op 接入 CP Core：

small-op attention backend
fused shared-KV attention backend

这也应该通过 adapter 或 attention backend 抽象接入，而不是让 CP Core 直接写 MF/Torch 特化逻辑。

12. 最终结论

HP 要同时支持 MF 和 Torch，DeepSeek V4 CP 的组件化边界应该是：

DeepSeekV4 CP Core:
    负责两阶段 CP 算法和显式通信。

MindFormers Adapter:
    负责把 MF CSA module 映射到统一协议。

Torch Adapter:
    负责把 Torch module/functions 映射到统一协议。

一句话：

不要让 CP Core 依赖 MF parent。
让 MF parent 只存在于 MindFormers Adapter 里。