Open

Hyper-Parallel 支持 MindFormers DeepSeek V4 CP 并行方案 #2

yide12Add on Jun 4

yide12Collaborators

Jun 4 Create

Hyper-Parallel 支持 MindFormers DeepSeek V4 CP 并行方案

背景与结论

MindFormers PR 8311 已实现 DeepSeek V4 hybrid attention 的基础网络和单卡 fused 算子调用路径。核心代码在：

mindformers/pynative/transformers/experimental_attention_variant/deepseek_v4_hybrid_attention.py
mindformers/pynative/transformers/experimental_attention_variant/csa.py

Hyper-Parallel PR 717 已补齐 DeepSeek V4 单卡 fused 路径依赖的三个 MindSpore 自定义算子：

npu_lightning_indexer_v2
npu_sparse_attn_sharedkv
npu_sparse_lightning_indexer_grad_kl_loss_v2

因此 HP 支持 DeepSeek V4 CP 的正确边界不是再实现一套 DeepSeek V4 模型，也不是把底层通信原语零散暴露给 MF，而是：

MF 提供 DeepSeek V4 基础网络结构、单卡 forward 语义、Compressor/Indexer/SharedKV attention 调用路径；HP 提供算子的分布式语义，以及一个完整的 DeepSeek V4 两段式 CP 功能接口。MF 只调用这个功能接口，不直接调用 HP 的 P2P、all-gather、select-and-pad 等通信原语。

职责边界

MF 负责

MF 保持 DeepSeek V4 网络本体所有权：

DeepSeekV4HybridAttention
CompressedSparseAttention
Compressor
CSAIndexer
FusedDSAIndexerLoss
MTP layer 的网络构造和 layer_number 映射
单卡 construct 路径，包括 SBND 到 BSND 的 layout 转换
csa_compress_ratios 的 per-layer ratio 分发
force_unfused_dsa debug path
基础网络的 FSDP/EP/单卡 fused/naive 对齐测试

HP 不新增 hyper_parallel/models/deepseek_v4/ 模型目录，不维护 DeepSeek V4 block、state_dict adapter 或 checkpoint 结构。

HP 负责

HP 的主要工作分两类。

第一类是算子的分布式语义：

为 PR 717 三个算子补 DistributedOp / layout inference / expand 语义。
明确这些算子在 CP 开启和不开启时的输入输出布局。
保证 MF 单卡 fused path 在 CP=1 时仍可直接调用 HP experimental ops。
保证 CP>1 时算子不会被普通 attention CP 规则误切分。

第二类是完整两段式 CP 能力：

HP 对 MF 暴露类似 torchtitan converter 的“压缩注意力前向替换”接口。
HP 在前向替换对象内部封装 window KV halo、ratio4 compressor halo、compressed KV all-gather、select-and-pad、global top-k、KL loss all-reduce 等细节。
MF 不直接调用这些通信原语，只把 attention module、输入张量和 CP meta 交给 HP。
HP 不遍历 MF 模型、不根据 MF 网络结构猜 hook 位置、不决定在哪个模块替换前向。
MF 在自己的 parallelize 流程里决定哪些输入边界需要 sequence shard、哪些 CSA module 需要替换前向为 HP 的 CP forward。
HP 管理 CP metadata：rank、local sequence 范围、packed sequence 的 cu_seqlens、compressed cu_seqlens。
HP 提供 CP=1 与 CP>1 的数值对齐测试。

MF 与 HP 的接口契约

MF 推荐暴露一个稳定 CP spec，避免 HP 通过易碎的私有属性猜模块结构：

class DeepSeekV4CPSpec:
    compress_ratio: int          # 0 / 4 / 128
    layout: str                  # BSND / TND
    window_size: int             # default 128
    ori_win_left: int            # default 127
    ori_win_right: int           # default 0
    softmax_scale: float
    has_compressor: bool
    has_indexer: bool
    force_unfused_dsa: bool

MF module 可以提供：

def get_deepseek_v4_cp_spec(self) -> DeepSeekV4CPSpec:
    ...

HP 提供两类接口，MF 在自己的 parallelize 流程中决定挂载位置。

MF 侧可以有一个自己的入口，例如：

def apply_deepseek_v4_context_parallel(model, cp_mesh, config):
    ...

但这个函数应属于 MF 的 DeepSeek V4 parallelize 逻辑。HP 不应该提供一个接收完整 MF model 并自行遍历模型结构的 apply(model, ...)。

第一类接口是类似 torchtitan CP patch 的压缩流程前向替换接口。torchtitan-npu 里有两种相关形态：SFA converter 用 NpuSparseAttention(parent) 这类 wrapper 替换模块；DeepSeek V4 CP patch 则直接替换 Attention.forward。HP 这里更应对齐后者：只替换 CompressedSparseAttention.construct 或注入 forward delegate，保留原 Cell 本体、参数注册和 state_dict 结构：

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __init__(self, origin_csa_module, cp_mesh, config):
        # 持有 MF 原始 CompressedSparseAttention。
        # 使用原 module 的 compressor / indexer / fused_indexer_loss / attn_sink 等权重和属性。
        ...

    def __call__(self, query, key, value, attention_mask, x=None, qr=None):
        # 等价替换 MF CompressedSparseAttention.construct 的 CP>1 前向路径。
        # 内部完整完成 Stage-1 + Stage-2。
        ...


def patch_deepseek_v4_compressed_sparse_attention_forward(
    csa_module,
    *,
    cp_mesh,
    config,
) -> object:
    # 返回原 csa_module；只 patch/委托 construct，不替换 module 本体。
    ...

MF 使用方式：

def apply_deepseek_v4_context_parallel_in_mf(model, cp_mesh, config):
    for layer in model.layers:
        # MF 自己知道 core_attention 是否是 DeepSeek V4 CSA。
        hp.patch_deepseek_v4_compressed_sparse_attention_forward(
            layer.attention.core_attention,
            cp_mesh=cp_mesh,
            config=config,
        )

    return model

这个接口形式和 torchtitan-npu 的 DeepSeek V4 CP patch 更接近：MF/模型侧选中目标 attention，HP 替换其前向流程。它也借鉴了 NpuSparseAttention(parent) / NpuLiCompute(parent) / NpuLiLoss(parent) 的“保留原模块能力、替换计算路径”思路，但实现上不替换整个 Cell 对象。HP 不扫描完整 model。

第二类接口是每个自定义算子的 DistributedOp / layout 语义。注意：这不是 DeepSeek V4 CP forward 主路径的编排接口，也不是通过 ContextParallel hook 自动改写输入。完整 CP attention forward 仍由 DeepSeekV4CompressedSparseAttentionCPForward 编排，但它在调用 PR717 op 前，需要把关键输入包装成 DTensor 或携带明确 layout/meta 的对象，从而触发对应 DistributedOp 语义处理多卡细节，例如 actual_seq_len、cu_seqlens_*、global/local sequence offset、loss reduce 等。

class ParallelNpuLightningIndexerV2(DistributedOp):
    def infer_layout(...):
        ...


class ParallelNpuSparseAttnSharedKV(DistributedOp):
    def infer_layout(...):
        ...


class ParallelNpuSparseLightningIndexerGradKLLossV2(DistributedOp):
    def infer_layout(...):
        ...

这些 DistributedOp 规则用于：

DeepSeek V4 CP forward 内部调用 PR717 op 时，通过 DTensor dispatcher 触发正确的 distributed op expand/layout 规则。
CP=1 或非 DeepSeek V4 CP forward 场景下，PR717 op 被 DTensor dispatcher 直接调用。
后续 MF 如果把 Op-1 / Op-2 / Op-3 暴露为独立分布式 op boundary，可以复用相同 layout 语义。
UT/ST 中验证这些 custom ops 在 DTensor 输入下不会被普通 attention CP 规则误切。

DeepSeek V4 CP 主路径不依赖 ContextParallel hook 来完成两段式 CP，但依赖这些 DistributedOp 语义来让 PR717 op 在多卡输入布局下正确处理 shape/meta/actual_seq_len。

前向替换接口语义上等价于 MF CompressedSparseAttention.construct(...) 的 CP 版本：

输入：MF attention module + 本地 query/key/hidden_states/qr + CP meta
输出：本地 sequence shard 的 attention output
内部：HP 自动完成 Stage-1 和 Stage-2

MF 不应该这样调用 HP：

halo = hp.exchange_window_halo(...)
cmp = hp.all_gather_compressed_kv(...)
idx = hp.select_and_pad(...)

这些函数即使在 HP 内部存在，也只作为 private implementation details。

风格对齐

“哪些 CompressedSparseAttention module 的前向被 HP CP forward 替换由 MF 决定”这个设计更接近 torchtitan 的 model-specific CP patch / converter 风格，而不是 MindSpeed 的全局 patch 风格。

对齐点：

像 torchtitan CP patch：模型侧知道自己的层结构和 attention 类型，因此由模型侧 parallelize/patch 选择目标 module，并将其 forward/construct 替换为专用 CP forward。
像 MindSpeed：输入 sequence 切分、position_ids、actual_seq_len / packed meta 这类数据路径处理，需要在 attention 前完成，不能等底层 attention op 自己猜。
不像 MindSpeed 的地方：HP 不通过全局 patch manager 按函数名或类名批量替换 MF 实现，也不根据 args 在 HP 内部扫描 MF 模型。

因此推荐表述是：

模块选择和 hook 位置：MF 所有，风格接近 torchtitan model-specific parallelize。
输入 sequence/meta 处理：借鉴 MindSpeed batch/context-parallel 思路。
两段式 attention CP 能力：HP 所有，通过 construct/forward 替换提供完整功能接口。

前向替换机制

DeepSeek V4 CP 仍建议沿用 HP 既有 CP 的 hook/patch 思路，但 hook 位置必须由 MF 决定。HP 不能用通用 attention CP hook 盲扫 MF 模型，也不能依赖 MF 网络内部命名去判断哪里是 DeepSeek V4 CSA。

MF 的 DeepSeek V4 parallelize 逻辑应安装两类能力：

用 HP 的 CP forward 替换 MF 的 CompressedSparseAttention.construct 压缩注意力流程。
确保进入该 CP forward 前，输入已经是本地 sequence shard，并提供必要 CP meta。

伪代码：

def mf_parallelize_deepseek_v4_context_parallel(model, cp_mesh, config):
    # 1. MF 自己选择需要替换的 CSA module。
    for csa_module in model.iter_deepseek_v4_csa_modules():
        hp.patch_deepseek_v4_compressed_sparse_attention_forward(
            csa_module,
            cp_mesh=cp_mesh,
            config=config,
        )

    return model

首版可落地机制是 HP 提供一个前向替换对象，对 MF 显式传入的 CompressedSparseAttention patch construct：

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __init__(self, parent, cp_mesh, config):
        self.origin_construct = parent.construct
        self.parent = parent
        self._hp_cp_mesh = cp_mesh
        self._hp_cp_config = config

    def __call__(self, query, key, value, attention_mask, x=None, qr=None):
        if not self._hp_cp_config.enabled:
            return self.origin_construct(query, key, value, attention_mask, x=x, qr=qr)

        return self._construct_cp_two_stage(query, key, value, attention_mask, x=x, qr=qr)


def patch_deepseek_v4_compressed_sparse_attention_forward(module, *, cp_mesh, config):
    cp_forward = DeepSeekV4CompressedSparseAttentionCPForward(module, cp_mesh, config)
    module._hp_origin_construct = module.construct
    module._hp_cp_forward = cp_forward

    def construct(query, key, value, attention_mask, x=None, qr=None):
        return cp_forward(query, key, value, attention_mask, x=x, qr=qr)

    module.construct = construct
    return module

这种方式继承了 torchtitan 的 replacement wrapper 思路，但实现上只替换前向，替换边界由 MF 选择。HP 只提供 DeepSeek V4 CSA 语义边界上的 CP forward，不理解也不依赖 MF 的完整模型结构。

替换前向不是包一层原 forward

这里的“patch construct”不能理解成简单包裹原 CompressedSparseAttention.construct。CP>1 时必须像 torchtitan 的 Attention.forward = attention_forward_with_cp 一样，在选定边界上重写 CP forward 逻辑。

推荐语义：

def construct(query, key, value, attention_mask, x=None, qr=None):
    if cp_enabled:
        # 完整 CP forward：Stage-1 + Stage-2。
        # 不能先调用原 construct，因为原 construct 会先本地 compressor/indexer。
        return cp_forward(query, key, value, attention_mask, x=x, qr=qr)

    return origin_construct(query, key, value, attention_mask, x=x, qr=qr)

不推荐语义：

def construct(...):
    out = origin_construct(...)   # 错：ratio4/128 已经在本地 compressed KV 上算错
    return hp_post_process(out)

因此 HP 的前向替换和 titan 的本质是一致的：都重写目标 attention 边界的 CP forward。区别只在实现载体：

titan DeepSeek V4 CP：替换 Attention.forward。
HP/MF 推荐：patch MF 选定 CompressedSparseAttention.construct，或在 MF 预留 delegate 后走 CP delegate。

如果 CompressedSparseAttention 边界拿不到 CP 必需信息，例如全局/per-doc RoPE 位置、local sequence 范围、packed meta，MF 应选择更高一层的 DeepSeekV4HybridAttention 作为 hook 边界。HP 不替 MF 判断边界，但 HP CP forward 要在入口处 loud fail，避免在信息不足时静默走错。

HP 是否要完整重写 attention

HP 需要实现一套完整的 DeepSeek V4 CP attention forward 编排，但不需要实现一套完整的 DeepSeek V4 网络本体。

需要 HP 完整负责：

CP 输入校验和 CP meta 消费。
ratio0 的 window KV halo + shared-KV attention。
ratio128 的 compressed KV all-gather + causal visible crop + shared-KV attention。
ratio4 的 compressor halo、global compressed KV、select-and-pad、global top-k、shared-KV attention、KL loss all-reduce。
反向通信语义，例如 halo backward、all-gather backward reduce-scatter、select-and-pad backward scatter-add。
调用 PR717 三个 custom ops 时的分布式 layout 规则。

不需要 HP 重新实现：

DeepSeek V4 block / layer / MTP 结构。
Q/K/QR projection 的参数定义和权重管理。
Compressor / CSAIndexer / FusedDSAIndexerLoss 的网络结构。
state_dict / checkpoint / config 分发。
MF 单卡 naive/fused reference path。

也就是说，HP 的 CP forward 不是：

重新写一个 DeepSeekV4HybridAttention

而是：

在 MF 选定的 attention 边界上，接管 CP>1 的注意力计算流程；
内部复用 MF module 的 compressor/indexer/attn_sink/softmax_scale 等成员，
但不调用原 construct 的本地 compressor/indexer 流程。

如果 HP 只提供三个 op 的 DistributedOp/layout 语义，而不提供这个完整 CP forward 编排，ratio4/128 会出问题：原 CompressedSparseAttention.construct 会先在本地 shard 上生成 compressed KV 和 top-k，错过 Stage-1 全局 compressed KV 构建，底层 op 再怎么补 layout 也已经晚了。

CP forward 内部仍要触发 DistributedOp

HP CP forward 不使用 ContextParallel hook 自动处理 DTensor，但在写前向时需要主动把将要进入 PR717 op 的张量包装成 DTensor 或明确 layout/meta，让 custom op 走自己的 DistributedOp 语义。

示意：

def _run_lightning_indexer_v2(self, q_indexer_local, global_cmp_kv, weights_local, cp_ctx):
    q_dt = DTensor.from_local(q_indexer_local, cp_ctx.mesh, (Shard(cp_ctx.seq_dim),))
    k_dt = DTensor.from_local(global_cmp_kv, cp_ctx.mesh, (Replicate(),))
    w_dt = DTensor.from_local(weights_local, cp_ctx.mesh, (Shard(cp_ctx.seq_dim),))

    return npu_lightning_indexer_v2(
        q_dt,
        k_dt,
        w_dt,
        actual_seq_qlen=cp_ctx.local_cu_seqlens_q,
        actual_seq_klen=cp_ctx.cu_seqlens_cmp_kv,
        layout_query=cp_ctx.layout,
        layout_key=cp_ctx.layout,
        cmp_ratio=cp_ctx.compress_ratio,
        ...
    )

类似地，调用 npu_sparse_attn_sharedkv 前：

q_dt = DTensor.from_local(q_local, mesh, (Shard(seq_dim),))
ori_kv_dt = wrap_prepared_window_kv(window_kv_with_halo, cp_ctx)
cmp_kv_dt = DTensor.from_local(global_cmp_kv, mesh, (Replicate(),))
idx_dt = DTensor.from_local(global_topk_indices, mesh, (Shard(seq_dim),))

out = npu_sparse_attn_sharedkv(
    q_dt,
    ori_kv=ori_kv_dt,
    cmp_kv=cmp_kv_dt,
    cmp_sparse_indices=idx_dt,
    cu_seqlens_q=cp_ctx.local_cu_seqlens_q,
    cu_seqlens_ori_kv=cp_ctx.local_cu_seqlens_ori_kv,
    cu_seqlens_cmp_kv=cp_ctx.cu_seqlens_cmp_kv,
    ...
)

这里 wrap_prepared_window_kv 的语义要非常明确：它不是普通 full-KV replicate，而是 HP CP forward 已经拼好 left halo 的 window KV。DistributedOp 规则需要知道这个输入已经是 prepared local window KV，不能再按普通 attention KV all-gather。

因此 DeepSeek V4 CP 主路径是：

HP CP forward 负责两段式算法和通信编排
        +
PR717 DistributedOp 负责 op 级 layout/meta/actual_seq_len 语义

输入 Sequence Sharding

MF 的 CP parallelize 不能只替换 CompressedSparseAttention 的前向。进入 attention 前，query、key、hidden_states(x)、low_rank_query(qr) 必须已经是本 rank 的 sequence shard，否则 HP CP forward 会把全局张量当本地张量处理，产生错误的 halo、compressed KV 和 top-k index。

输入 sequence sharding 不再单独抽成 DeepSeekV4SequenceShardStyle。更好的形态是：

MF 在自己的输入/batch/layer boundary 完成 sequence sharding 和 DeepSeekV4CPContext 构建。
HP 的 CP forward 在 construct 开始处强校验 query/key/x/qr 是否已经是本地 sequence shard。

MF PyNative 当前 CompressedSparseAttention.construct 的内部张量是 seq-first：

query: [sq_local, b, n_head, v_head_dim]
key:   [sq_local, b, 1,      v_head_dim]
x:     [sq_local, b, hidden]
qr:    [sq_local, b, q_lora_rank]

因此 HP 在 MF attention 边界的 sequence dim 是 0。等到调用 npu_sparse_attn_sharedkv 前，HP CP forward 再按算子要求转为 BSND，此时 op 侧 sequence dim 是 1。这两个 layout 语义要在 CP forward 里固定下来，不能让 MF 手动处理。

DeepSeekV4CompressedSparseAttentionCPForward.__call__(...) 入口应做强校验：

def _validate_local_sequence_shard(query, key, hidden_states, low_rank_query, cp_ctx):
    assert is_seq_sharded_or_local(query, seq_dim=0, cp_ctx=cp_ctx)
    assert is_seq_sharded_or_local(key, seq_dim=0, cp_ctx=cp_ctx)
    assert hidden_states is None or is_seq_sharded_or_local(hidden_states, seq_dim=0, cp_ctx=cp_ctx)
    assert low_rank_query is None or is_seq_sharded_or_local(low_rank_query, seq_dim=0, cp_ctx=cp_ctx)

如果 MF 的数据路径已经提前把 batch 切成 local sequence shard，HP CP forward 只做校验。如果 MF 传入的是全局 sequence tensor，应由 MF 在自己的输入/batch/layer boundary 调用 HP 的通用 ContextParallel style 或本地切分逻辑先转成 Shard(seq) DTensor；不应等到 CompressedSparseAttention 内部再临时切分。

HP 侧新增模块建议

建议新增：

hyper_parallel/core/context_parallel/deepseek_v4_context_parallel.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_flow.py
hyper_parallel/core/context_parallel/npu_lightning_indexer_v2_context_parallel.py
hyper_parallel/core/context_parallel/npu_sparse_attn_sharedkv_context_parallel.py
hyper_parallel/core/context_parallel/npu_sparse_lightning_indexer_grad_kl_loss_v2_context_parallel.py
hyper_parallel/core/shard/ops/parallel_npu_lightning_indexer_v2.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_attn_sharedkv.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py
hyper_parallel/core/shard/ops/yaml/npu_lightning_indexer_v2_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_attn_sharedkv_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_lightning_indexer_grad_kl_loss_v2_ops.yaml

核心数据结构：

class DeepSeekV4CPContext:
    cp_rank: int
    cp_size: int
    cp_group: object
    global_seq_len: int
    local_seq_start: int
    local_seq_end: int
    local_seq_len: int
    compress_ratio: int
    layout: str
    window_size: int
    cu_seqlens_q: Tensor | None
    cu_seqlens_ori_kv: Tensor | None
    cu_seqlens_cmp_kv: Tensor | None
    cu_seqlens_padded: Tensor | None
    cu_seqlens_real: Tensor | None

前向替换对象：

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __init__(self, origin_csa_module, cp_mesh, config):
        ...

    def __call__(self, query, key, value, attention_mask, x=None, qr=None):
        ...

    # private/internal helpers, not exposed to MF
    def _stage1_build_global_compressed_kv(self, csa_module, hidden_states, cp_ctx):
        ...

    def _stage2_run_local_sharedkv_attention(
        self, csa_module, query, ori_kv, global_compressed_kv, cp_ctx
    ):
        ...

__call__ 是替换 MF CompressedSparseAttention.construct 的完整 CP forward。_stage1_*、_stage2_* 和所有通信实现只在 HP 内部使用。

不建议对 MF 暴露如下接口：

class DeepSeekV4CPRuntime:
    # 不推荐：这会把两段式 CP 拆散给 MF 编排
    def build_compressed_kv(self, csa_module, hidden_states, cp_ctx):
        ...

    # 不推荐：这要求 MF 理解 HP 的 Stage-1 产物和通信细节
    def run_sharedkv_attention(self, csa_module, query, ori_kv, compressed_kv, cp_ctx):
        ...

二阶段 CP 总体算法

DeepSeek V4 不能套普通 ring/Ulysses CP。HP 需要实现 MF RFC 中的二阶段算法，并把它封装成一个前向替换对象：

Stage-1: 构建全局 compressed KV
Stage-2: 本地 Q 调 shared-KV sparse attention

Stage-1 只处理 compressed KV 的全局一致性。Stage-2 中 Q 不 all-gather，每个 rank 只输出自己的 local sequence shard。

MF 视角：

hp.patch_deepseek_v4_compressed_sparse_attention_forward(
    origin_csa,
    cp_mesh=cp_mesh,
    config=config,
)
output = origin_csa(query, key, value, attention_mask, x=x, qr=qr)

HP 内部视角：

1. 读取 csa_module 的 ratio/window/scale/sink/compressor/indexer/loss 配置
2. Stage-1 构建 global compressed KV
3. Stage-2 构建 window ori_kv 视图并调用 shared-KV attention
4. ratio4 训练态自动处理 indexer top-k 和 KL loss all-reduce
5. 返回与 MF 单卡 construct 相同 layout 的本地 output

Stage-1：构建 compressed KV

ratio=0

纯滑窗 attention，无 compressed KV：

compressed_kv = None
cmp_sparse_indices = None

HP CP forward 内部为 Stage-2 准备 window KV 左 halo，对 MF 不暴露 halo 接口。

ratio=128

HCA，非重叠压缩：

1. MF compressor(local hidden_states) -> local compressed_kv
2. HP all-gather 所有 rank 的 local compressed_kv
3. HP 按因果可见范围和 packed sequence meta 裁剪/组织 global compressed_kv
4. 重建 cu_seqlens_cmp_kv

首版要求：

T_pad % (cp_size * 128) == 0

packed sequence 下不能简单使用 cu_seqlens_padded // 128，而应按每个文档的 floor(real_len / 128) 构建 compressed cu_seqlens。

ratio=4

CSA，重叠压缩，最复杂。HP 需要实现：

1. rank r 将本地最后 ratio 个 token 发送给 rank r+1
2. rank r 从 rank r-1 接收 ratio-token compressor halo
3. MF compressor 在 [halo, local] 上生成 s / ratio + 1 个 compressed slots
4. slot0 统一无效：
   - rank0: 前导 PAD
   - rank>0: 边界 DUP
5. all-gather 所有 rank 的等长 compressed slots
6. select-and-pad 去掉 PAD/DUP/半组/文档 padding 产生的无效块
7. 生成全局有序 compressed_kv 和 cu_seqlens_cmp_kv

首版 select-and-pad 可先用小算子实现：

cmp_padded = concat([cmp_gathered, zero_row])
cmp_full = index_select(cmp_padded, gather_index)

gather_index 只依赖 cu_seqlens_padded、cu_seqlens_real、cp_size、compress_ratio，可以 host 端预计算并缓存。后续再替换为 CANN 融合算子。

Stage-2：本地 Q shared-KV sparse attention

所有 ratio 都需要 window KV 左 halo：

rank r 发送 local ori_kv[-127:] 给 rank r+1
rank r 接收 rank r-1 的 127 个 ori_kv
ori_kv_for_attn = concat(left_halo, local_ori_kv)

rank0 的 left halo 用 0/pad 填充，并通过 mask 保证不可见。Backward 需要把 halo grad 反向传回左邻。

ratio=0

out, lse = npu_sparse_attn_sharedkv(
    q_local,
    ori_kv=window_kv_with_left_halo,
    cmp_kv=None,
    cmp_sparse_indices=None,
    sinks=csa_module.attn_sink,
    softmax_scale=csa_module.softmax_scale,
    cmp_ratio=None,
    ori_win_left=127,
)

ratio=128

out, lse = npu_sparse_attn_sharedkv(
    q_local,
    ori_kv=window_kv_with_left_halo,
    cmp_kv=global_compressed_kv,
    cmp_sparse_indices=None,
    sinks=csa_module.attn_sink,
    softmax_scale=csa_module.softmax_scale,
    cmp_ratio=128,
    ori_win_left=127,
)

ratio=4

先在全局 compressed KV 上运行 indexer：

topk_indices, topk_score = npu_lightning_indexer_v2(
    q_indexer_local,
    global_compressed_kv,
    weights_local,
    layout_query=layout,
    layout_key=layout,
    sparse_count=index_topk,
    sparse_mode=3,
    cmp_ratio=4,
    return_value=True,
)

再运行 shared-KV attention：

out, lse = npu_sparse_attn_sharedkv(
    q_local,
    ori_kv=window_kv_with_left_halo,
    cmp_kv=global_compressed_kv,
    cmp_sparse_indices=topk_indices,
    sinks=csa_module.attn_sink,
    softmax_scale=csa_module.softmax_scale,
    cmp_ratio=4,
    ori_win_left=127,
)

训练态 KL loss：

loss = npu_sparse_lightning_indexer_grad_kl_loss_v2(...)
loss = all_reduce(loss, group=cp_group)

如果 MF 继续在 CompressedSparseAttention.construct 内部调用 indexer/loss，MF 应在 parallelize 阶段对该 module 调用 patch_deepseek_v4_compressed_sparse_attention_forward(...)：CP>1 时由 HP CP forward 完整接管 Stage-1/Stage-2。MF 不应接收 Stage-1 产物再自行编排 Stage-2，否则会把 HP 两段式能力拆散，接口也会变得不可维护。

三个自定义算子的分布式语义

npu_lightning_indexer_v2

输入布局：

query: 本地 sequence shard，BSND/TND。
key: Stage-1 得到的全局 compressed KV。
weights: 本地 sequence shard。

输出：

sparse_indices: 本地 Q 对全局 compressed KV 的索引。
sparse_values: 本地 Q 的 top-k score。

CP 规则：

query/weights 沿 sequence 维 shard。
key 默认 replicate/global。
TND 下 actual_seq_qlen 是本地 Q cu_seqlens，actual_seq_klen 是 global compressed cu_seqlens。
ratio4 必须传 cmp_ratio=4。

npu_sparse_attn_sharedkv

输入布局：

q: 本地 sequence shard。
ori_kv: 本地 ori_kv + left halo。
cmp_kv: Stage-1 全局 compressed KV，ratio0 为 None。
cmp_sparse_indices: ratio4 的全局 compressed KV index，ratio0/128 为 None。

输出：

attention_out: 继承本地 Q sequence shard。
softmax_lse: 本地 Q 对应 layout。

CP 规则：

不允许把 ori_kv 当普通 Replicate KV 自动 all-gather。
ori_kv 的 halo 通信由 DeepSeekV4 CP forward 完成。
cmp_kv 由 Stage-1 CP forward 构建。
TND 下显式传 cu_seqlens_q、cu_seqlens_ori_kv、cu_seqlens_cmp_kv。

npu_sparse_lightning_indexer_grad_kl_loss_v2

输入布局：

query: 本地 main attention Q。
key: global compressed KV 或其 teacher key。
query_index / weights: 本地 indexer branch。
key_index: global compressed index key。
sparse_indices: 本地 Q 指向 global compressed KV 的 index。

CP 规则：

loss 标量跨 CP group all-reduce。
d_query_index / d_weights 留在本地 shard。
d_key_index / d_cmp_kv 按 Stage-1 compressed KV 分布规则回传。
softmax_max/softmax_sum 为 None 的 ratio4 模式必须维持。

Packed Sequence / TND 元数据

HP 不负责 MF dataloader，但 HP CP forward 和 per-op DistributedOp/layout 规则必须消费 MF 传入的 packed sequence meta。

关键规则：

RoPE 按文档内 position reset，不能简单使用全局 arange。
compressed pooling 不能跨文档。
每个文档 padded length 必须满足 doc_padded_len % ratio == 0。
compressed cu_seqlens 使用每个文档 floor(real_len / ratio) 的 prefix sum。
T_pad % (cp_size * ratio) == 0 不满足时应 loud fail。

实施阶段

Phase 1：算子接入和分布式注册

合入/对齐 PR 717 的三个 MindSpore 自定义算子。
补 npu_lightning_indexer_v2 distributed op。
补 npu_sparse_attn_sharedkv distributed op。
补 npu_sparse_lightning_indexer_grad_kl_loss_v2 distributed op。
保证 MF 8311 单卡 fused path 可直接 import HP experimental ops。

Phase 2：MF/HP 能力接口

MF 暴露 get_deepseek_v4_cp_spec() 或等价稳定属性。
HP 交付 patch_deepseek_v4_compressed_sparse_attention_forward(csa_module, cp_mesh, config)。
HP 交付 DeepSeekV4CompressedSparseAttentionCPForward，替换 MF 显式传入 CSA module 的 construct。
HP 交付 ParallelNpuLightningIndexerV2。
HP 交付 ParallelNpuSparseAttnSharedKV。
HP 交付 ParallelNpuSparseLightningIndexerGradKLLossV2。
MF 自己决定在哪个输入边界挂 sequence-shard hook，以及哪些 CompressedSparseAttention patch 为 HP CP forward。
MF 只传 module、输入张量和 CP meta，不让 HP 反向依赖 MF 网络结构，也不编排 Stage-1/Stage-2 的通信细节。

Phase 3：ratio0 / ratio128

在 HP CP forward 内部实现 ratio0 的 window attention CP。
在 HP CP forward 内部实现 ratio128 的 HCA 两段式 CP。
支持 BSND。
CP=2/4 对齐 CP=1。

Phase 4：ratio4 CSA 两段式前半段

在 HP CP forward 内部完成 ratio-token compressor halo。
在 HP CP forward 内部完成等长 compressed slots all-gather。
在 HP CP forward 内部完成 select-and-pad 小算子。
在 HP CP forward 内部重建 cu_seqlens_cmp_kv。
对 MF 仍保持“patch CSA construct 后正常调用 construct”的接口。

Phase 5：ratio4 CSA 两段式后半段

在 HP CP forward 内部让 indexer 在 global compressed KV 上 top-k。
在 HP CP forward 内部让 shared-KV attention 消费 global cmp_sparse_indices。
在 HP CP forward 内部完成 KL loss CP all-reduce。
backward 对齐。

Phase 6：TND / packed sequence

per-doc RoPE metadata。
不跨文档压缩。
半组丢弃。
多文档 CP 对齐测试。

Phase 7：TP+CP 和性能优化

TP>1 + CP 组合。
select-and-pad 融合算子。
通信 overlap。
zigzag/2xCP 负载均衡后续再做。

最小测试矩阵

PR 717 三个算子单卡 forward/backward UT/ST。
MF 8311 单卡 ratio0/4/128 fused path 能 import HP ops 并运行。
MF 侧 apply_deepseek_v4_context_parallel 调用 HP 能力后，选定输入边界具备 sequence-shard hook，选定 CSA module 的 construct 被 patch 为 HP CP forward。
CompressedSparseAttention 入口处 query/key/x/qr 已是 Shard(seq_dim=0) DTensor 或本地 sequence shard。
ratio0 CP=2，window halo 后输出对齐 CP=1。
ratio128 CP=2，compressed KV all-gather 后输出对齐 CP=1。
ratio4 Stage-1，select-and-pad 后 compressed KV 对齐 CP=1。
ratio4 Stage-2，top-k indices 指向全局 compressed KV。
ratio4 KL loss all-reduce 后 loss tracker 对齐 CP=1。
TND 两文档跨 rank，RoPE reset 和 cu_seqlens 正确。
T_pad % (cp_size * ratio) != 0 loud fail。
doc_padded_len % ratio != 0 loud fail。

风险点

MF 当前 CompressedSparseAttention.construct 会在本地直接调用 self.compressor(x)、self.indexer(...) 和 npu_sparse_attn_sharedkv(...)。CP>1 不能只给底层 op 注册普通 DTensor 规则，否则 ratio4/128 会在本地 compressed KV 上产生静默错误。
npu_sparse_attn_sharedkv 的 ori_kv 是 window KV + halo，不是普通 full KV replicate。
ratio4 的 cmp_sparse_indices 必须是全局 compressed KV 索引，不能是每 rank 本地 offset。
TND 下 compressed cu_seqlens 不能从 padded cu_seqlens 直接整除得到。
KL loss 是辅助 loss，CP 下不 all-reduce 会造成训练日志和梯度缩放语义偏差。

一句话方案

HP 支持 DeepSeek V4 CP 的方案是：MF 继续拥有 DeepSeek V4 基础网络；HP 主要交付两件事：一是 PR 717 三个 custom ops 的 DistributedOp/layout 语义，二是类似 torchtitan 的 DeepSeekV4CompressedSparseAttentionCPForward 前向替换对象。MF 选择 CSA module 并 patch 其 construct；HP CP forward 负责编排 window halo、compressed KV 构建、select-and-pad、global top-k、shared-KV attention 和 KL loss all-reduce，并在调用 PR717 op 前主动包装 DTensor/传入 meta，触发对应 DistributedOp 处理 actual_seq_len、cu_seqlens 和 layout。

yide12Collaborators

Jun 5 Comments：

MindFormers 两段式 DeepSeek V4 CP 实现拆解

本文只关注“完整两段式 CP 能力”这一类问题：

MF 如何选择 attention 边界、注入 CP meta，并在 CP forward 内封装 window KV halo、ratio4 compressor halo、compressed KV all-gather、select-and-pad、global top-k、KL loss 等细节。

不展开 PR717 custom op 的 DistributedOp/layout 语义。

参考实现：

/home/wyd/code_review_deepseek_v4_cp/mindformers_pr8311_latest
PR8311 head: a74de4127e8e72ae633ca21c3605e4a956d34399
主要文件：
- mindformers/pynative/base_models/gpt/parallelize.py
- mindformers/pynative/distributed/context_parallel.py
- mindformers/pynative/transformers/experimental_attention_variant/csa.py
- mindformers/pynative/transformers/experimental_attention_variant/csa_context_parallel.py

1. 先给结论

MF 当前没有调用一个来自 HP 的 DeepSeekV4CompressedSparseAttentionCPForward 对象。它是在 MF 仓内实现了同构能力：

MF parallelize 选择 CompressedSparseAttention module
  -> 挂一个 light forward-pre-hook
  -> hook 只注入 _cp_ctx = (cp_size, cp_rank, cp_group)
  -> CompressedSparseAttention.construct 发现 _cp_ctx 后委托给 two_stage_cp_naive_attention
  -> two_stage_cp_naive_attention 封装完整 Stage-1 + Stage-2

所以从接口职责上看，它已经验证了我们想让 HP 提供的“压缩注意力前向替换”边界；只是当前实现位置在 MF，且走 local Tensor + naive small-op，不是 HP fused/DTensor 版本。

flowchart TB
    A["parallelize_gptmodel"] --> B["Phase 3: apply_context_parallel_attention"]
    B --> C{"core_attention 是 CompressedSparseAttention?"}
    C -->|否| D["普通 attention<br/>套已有 ContextParallelAttentionStyle"]
    C -->|是| E["CompressedSparseAttentionContextParallel(cp_mesh)._apply"]
    E --> F["注册 forward-pre-hook"]
    F --> G["hook 注入 _cp_ctx"]
    G --> H["CompressedSparseAttention.construct"]
    H --> I{"_cp_ctx and cp_size > 1?"}
    I -->|否| J["原单卡路径"]
    I -->|是| K["two_stage_cp_naive_attention"]
    K --> L["Stage-1: global compressed KV"]
    L --> M["Stage-2: local Q attention"]

2. MF 如何决定替换哪个 forward

入口在 parallelize_gptmodel：

Phase 1 TP
Phase 2 EP
Phase 3 CP attention hook
Phase 4 recompute/swap
Phase 5 FSDP
Phase 6 CP root I/O hook

CP attention hook 的核心逻辑在 apply_context_parallel_attention：

遍历 gpt_model.decoder.layers。
在每个 layer 内查找 self_attention.*core_attention。
如果该 module 是 CompressedSparseAttention，安装 DeepSeek V4 专用 two-stage CP hook。
否则普通 attention 继续走原有 CP style。

这点和目标 HP 接口边界一致：

MF 自己遍历 MF 模型。
MF 自己知道哪个 module 是 DeepSeek V4 CSA/HCA attention。
HP 不应该扫描 MF 模型，也不应该按 MF 私有命名猜 hook 位置。

当前 MF 代码等价于未来这样的调用：

for layer in decoder.layers:
    csa = find_core_attention(layer)
    if isinstance(csa, CompressedSparseAttention):
        hp.patch_deepseek_v4_compressed_sparse_attention_forward(
            csa,
            cp_mesh=cp_mesh,
            config=config,
        )

区别只是：现在 patch 逻辑由 MF 自己的 CompressedSparseAttentionContextParallel 完成。

3. MF 如何传 CP meta

CompressedSparseAttentionContextParallel 初始化时从 cp_mesh 取三项：

cp_size = cp_mesh.size()
cp_rank = cp_mesh.get_local_rank()
cp_group = cp_mesh.get_group()

然后 _apply(module) 注册 pre-hook：

new_kwargs["_cp_ctx"] = (cp_size, cp_rank, cp_group)

这个 hook 有两个重要约束：

hook 不做 compute。
hook 不做 communication。

原因是通信必须发生在 construct 的 autograd 上下文里，否则容易遇到 hook 内 recompute/top-cell 相关问题，也不好保证反向 collective 对称。

sequenceDiagram
    participant P as MF parallelize
    participant M as CompressedSparseAttention
    participant H as CP pre-hook
    participant C as construct
    participant F as two_stage_cp_naive_attention

    P->>M: register_forward_pre_hook
    M->>H: forward call enters hook
    H->>H: build kwargs['_cp_ctx']
    H-->>C: no compute, no comm
    C->>C: check _cp_ctx
    C->>F: delegate heavy CP forward

未来 HP 化时，CP meta 应该比当前 MF 更完整。当前 _cp_ctx 只有 rank/size/group；packed sequence 所需的 cu_seqlens_padded / cu_seqlens_real / cmp_cu_seqlens 没有进入 runtime 主路径。

建议 HP 目标 meta 至少包含：

cp_size
cp_rank
cp_group
global_seq_len
local_seq_start
local_seq_len
position_ids_local
cu_seqlens_padded
cu_seqlens_real
cu_seqlens_cmp_kv

4. MF 如何保证输入已经是 local sequence

两段式 attention forward 不负责切 input_ids / labels / loss_mask / position_ids。MF 在 root model 边界另挂了 ContextParallelModelIOStyle。

它处理的输入名：

input_ids
position_ids
attention_mask
decoder_input
labels
loss_mask
actual_seq_len

普通 BS 路径：

[B, S] -> [B, S / CP]

TND/packed 路径检测到 actual_seq_len 后：

[B, S] -> reshape [1, B*S] -> [1, B*S / CP]

当前实现要求 position_ids 必须存在。main Q/K RoPE 依赖切分后的 local position_ids 保持全局位置语义。

flowchart TB
    A["全局 batch"] --> B["ContextParallelModelIOStyle"]
    B --> C["prepare_context_parallel_input"]
    C --> D["input_ids_local"]
    C --> E["labels_local"]
    C --> F["loss_mask_local"]
    C --> G["position_ids_local"]
    G --> H["GPTModel rotary_pos_emb(position_ids=...)"]
    H --> I["main Q/K 已按全局位置 RoPE"]
    I --> J["进入 CompressedSparseAttention 的 query/key/x/qr 已是 local shard"]

这也是未来 HP 接口要遵守的边界：HP 的 CP attention forward 可以校验输入 shape/meta，但不应该在 CompressedSparseAttention 内临时切全局 batch。

5. CP forward 的实际入口

CompressedSparseAttention.construct 增加 _cp_ctx=None 参数。

逻辑是：

if _cp_ctx is not None and cp_size > 1:
    if not force_unfused:
        raise ValueError
    return two_stage_cp_naive_attention(self, query, ori_kv, x, qr, *cp_ctx)
else:
    run original single-card construct

这等价于前向替换对象的 __call__：

def __call__(query, key, value, attention_mask, x=None, qr=None):
    if cp_disabled:
        return origin_construct(...)
    return construct_cp_two_stage(...)

当前限制：

CP>1 只支持 force_unfused_dsa=True。
fused npu_sparse_attn_sharedkv CP 主路径未实现。
因此 MF 当前是在 small-op naive 路径上验证算法，不是最终性能路径。

6. Stage-1：compressed KV 全局化

two_stage_cp_naive_attention 内部首先计算：

ratio = cell.compress_ratio
s_local = query.shape[0]
start = cp_rank * s_local
overlap = ratio == 4

6.1 ratio4 CSA：compressor halo

ratio4 是重叠压缩。边界 block 需要左邻最后 ratio 个 token。

当前实现没有用 P2P，而是：

x_bnd = all_gather(x[-ratio:])
halo_x = x_bnd[(cp_rank - 1) * ratio : cp_rank * ratio]

然后对 main compressor 和 indexer compressor 都执行：

local_slots = stage1_compress_with_halo(compressor, x, halo_x, cp_rank, overlap=True)
gathered = all_gather(local_slots)
gather_index = build_select_pad_gather_index(...)
cmp_global = select_and_pad_smallop(gathered, gather_index)

flowchart TB
    A["x local shard"] --> B["all_gather 每 rank x[-ratio:]"]
    B --> C["取左邻 halo_x"]
    C --> D1["cell.compressor<br/>stage1_compress_with_halo"]
    C --> D2["cell.indexer.compressor<br/>stage1_compress_with_halo"]
    D1 --> E1["all_gather local_slots"]
    D2 --> E2["all_gather local_slots"]
    E1 --> F1["select-and-pad"]
    E2 --> F2["select-and-pad"]
    F1 --> G1["cmp_main_global"]
    F2 --> G2["cmp_idx_global"]

slot 语义：

rank0：slot0 = PAD，后面是本 rank valid blocks。
rank>0：slot0 = boundary DUP，后面是本 rank valid blocks。
select-and-pad 删除 PAD/DUP/无效半组/padding block。

6.2 ratio128 HCA：非重叠 compressed KV all-gather

ratio128 不需要 halo_x：

local_slots = compressor(local_x, rope_pos_offset=rank * s)
gathered = all_gather(local_slots)
cmp_main_global = select_and_pad_smallop(gathered, gather_index)

HCA 没有 indexer，所以没有 cmp_idx_global。

6.3 当前 packed meta 缺口

build_select_pad_gather_index 支持：

cu_seqlens_padded
cu_seqlens_real
cp
ratio

但 runtime 当前调用是：

build_select_pad_gather_index(
    [0, s_local * cp_size],
    [0, s_local * cp_size],
    cp_size,
    ratio,
    overlap=overlap,
)

也就是只按单段 full-valid sequence 构造 gather index。packed/TND 的真实 cu_seqlens 语义有 helper 和 UT，但没有接进 two-stage CP runtime 主路径。

对 HP 来说，这是必须补齐的 CP metadata 管理点。

7. Stage-2：local Q attention

Stage-2 的原则：

Q 不 all-gather
输出不 scatter-back
每 rank 只算本地 query shard 的 output

7.1 window KV halo

对 window ori_kv：

hlen = min(window - 1, s_local)
ori_bnd = all_gather(ori_kv[-hlen:])
ori_ext, halo_len = assemble_window_halo(ori_bnd, ori_kv, cp_rank, hlen, window)

assemble_window_halo 支持 window - 1 > s_local，也就是 halo 可能跨多个左侧 rank，不只取直接左邻。

7.2 调用 `_construct_naive`

Stage-2 最后委托回 MF 已验证的 naive attention：

out = cell._construct_naive(
    query,
    ori_ext,
    x,
    qr,
    pos_offset=start,
    cmp_kv_global=cmp_main,
    ori_halo_len=halo_len,
    cmp_idx_global=cmp_idx,
)

_construct_naive 因 CP 新增的参数含义：

pos_offset: 本 rank local query 的全局起点。
cmp_kv_global: Stage-1 得到的全局 main compressed KV。
ori_halo_len: ori_ext 前面拼了多少 window halo。
cmp_idx_global: Stage-1 得到的全局 indexer compressed K。

flowchart TB
    A["local query"] --> H["cell._construct_naive"]
    B["ori_ext = left halo + local ori_kv"] --> H
    C["cmp_main_global"] --> H
    D["cmp_idx_global, ratio4 only"] --> H
    E["pos_offset = cp_rank * s_local"] --> H
    F["ori_halo_len"] --> H
    H --> I["local output shard"]

7.3 HCA global visibility

ratio128 没有 indexer。_construct_naive 调：

get_compress_topk_idxs(
    ratio,
    batch,
    local_seq_len,
    offset=len(ori_ext),
    pos_offset=global_start,
    n_compressed=len(cmp_main_global),
)

这样每个本地 query 的 compressed visibility 按全局位置计算。

7.4 CSA global top-k

ratio4 有 indexer。CP 下不能用本地 compressed K 做 top-k，必须用 cmp_idx_global：

q_idx, _, w_idx = indexer.forward_before_topk(
    x_local,
    qr_local,
    query_pos_offset=global_start,
)

topk = fused_qk_topk_naive(
    q_idx,
    cmp_idx_global,
    w_idx,
    eff_topk,
    global_causal_mask,
)

关键修正：

indexer Q 的 RoPE 用 query_pos_offset。
key 是 global indexer compressed K。
causal mask 用 global query positions。
top-k 宽度按 key length 截断，而不是 local query length。

8. KL loss 当前是怎么做的

ratio4 training 时，_construct_naive 会计算 naive KL：

compute_dsa_indexer_loss(...)
output = DSAIndexerLossAutoScaler(output, indexer_loss)
save_to_indexer_losses_tracker(indexer_loss, ...)

需要明确：MF 当前 two-stage CP forward 内没有显式做 CP all-reduce KL loss 标量。

源码注释给出的边界是：

naive compute_dsa_indexer_loss 不做内部 all-reduce。
tracker 是 single-card tracker。
cross-rank reduce 由 trainer/callback step end 负责。

因此你列的“KL loss all-reduce”更适合作为 HP 目标设计要求，而不是 MF 当前 CP forward 已实现事实。HP 迁移时应把 KL loss 分成两件事：

梯度路径：indexer loss 必须正确接入 output graph。
标量统计/scale：跨 CP rank 的 reduce、tracker 记录、loss scale 需要定义清楚。

flowchart TB
    A["ratio4 global top-k"] --> B["compute_dsa_indexer_loss"]
    B --> C["DSAIndexerLossAutoScaler<br/>接梯度"]
    B --> D["save_to_indexer_losses_tracker<br/>记录统计"]
    D --> E["trainer/callback<br/>外层 parallel loss 修正"]
    B -.当前 CP forward 内没有.-> F["显式 CP all-reduce KL"]

9. 反向通信对称性

MF 当前通信用函数式 all_gather_into_tensor。代码注释说明它的 backward 是 reduce-scatter VJP。

为了避免某些 gathered tensor 在某些 rank 上没有进入 loss graph，导致 backward collective 不对称，MF 做了 zero-weight tie：

for t in ties:
    out = out + ops.cast(t.sum(), out.dtype) * 0.0

被 tie 的 tensor 包括：

ratio4 x_bnd
window ori_bnd

这不是数值修正，而是 collective 对称性修正。HP 实现时必须保留同等语义；如果未来把 halo all-gather 换成 P2P，需要单独实现 autograd wrapper。

10. MF 测试如何证明 CP forward 正确

MF 的测试分三层。

flowchart TB
    A["helper / unit tests"] --> A1["gather_index END-based"]
    A --> A2["cmp_cu_seqlens"]
    A --> A3["select-and-pad"]

    B["single-process CP reconstruction"] --> B1["Stage-1 cmp == CP=1"]
    B --> B2["RoPE offset == CP=1"]
    B --> B3["ratio0 / ratio128 / ratio4 output == CP=1"]

    C["real multi-card ST"] --> C1["2-card Stage-1 real all-gather"]
    C --> C2["whole-net CP=2 smoke"]
    C --> C3["same checkpoint CP=1 vs CP=2/4 loss equivalence"]

    A --> B --> C

最关键的测试：

test_csa_context_parallel.py
- 验证 Stage-1、window halo、HCA、CSA、RoPE offset、top-k cap。
test_two_card_cp_comm.py
- 真实 2 卡通信验证 global compressed KV 对齐 CP=1。
test_cp_loss_equivalence.py
- 同一 checkpoint 下 CP=1 vs CP=2/4 逐 step loss 对齐。

这正对应 HP 未来需要提供的“CP=1 与 CP>1 数值对齐测试”。

11. 和你列出的 HP 第二类能力逐项对照

目标能力	MF 当前实现方式	HP 化建议
压缩注意力前向替换接口	`CompressedSparseAttentionContextParallel` + `_cp_ctx` + `two_stage_cp_naive_attention`	HP 提供 `DeepSeekV4CompressedSparseAttentionCPForward` 或 patch 函数，MF 负责选择 module
HP 不遍历 MF 模型	当前 MF 自己遍历 `decoder.layers` 并识别 `CompressedSparseAttention`	保持该边界，HP 只处理显式传入的 CSA module
MF 决定 sequence shard 边界	`ContextParallelModelIOStyle` 挂在 root model I/O	HP CP forward 不应临时切全局 batch，只消费 local shard 和 meta
CP metadata	当前只有 `_cp_ctx=(size, rank, group)`，local range 用 `rank * s_local` 推导	HP 需要扩展为显式 meta，特别是 packed `cu_seqlens`
window KV halo	`ori_kv[-hlen:]` all-gather + `assemble_window_halo`	HP 可先复用 all-gather 语义，后续优化 P2P
ratio4 compressor halo	`x[-ratio:]` all-gather，取左邻 halo	HP 需要支持 main compressor 和 indexer compressor 两套 Stage-1
compressed KV all-gather	`_ag_seq(local_slots, cp_group)`	HP fused/DTensor 版本仍要保证 backward reduce-scatter 对称
select-and-pad	host `gather_index` + `cat zero row + index_select`	HP 可先 small-op，后续 fused select-and-pad
global top-k	local indexer Q/weights + global `cmp_idx` + global causal mask	HP 不能在本地 compressed K 上做 top-k
KL loss all-reduce	当前 CP forward 内没有显式 all-reduce，依赖外层统计/scale	HP 需要明确 Op-3/naive loss 的 reduce 与 tracker 语义
CP=1/CP>1 对齐测试	已有 helper、single-process golden、2-card comm、whole-net loss equivalence	HP 应保留同样分层测试

12. 对 HP 接口的直接启示

未来 HP 接口建议长这样：

def patch_deepseek_v4_compressed_sparse_attention_forward(
    csa_module,
    *,
    cp_mesh,
    cp_meta_provider,
):
    ...

其中：

csa_module 由 MF 传入。
cp_mesh 提供 rank/size/group。
cp_meta_provider 提供 local/global sequence、position ids、packed cu_seqlens。
HP patch 后的 forward 内部完整封装 Stage-1 + Stage-2。
MF 不直接调用 exchange_halo / all_gather_cmp / select_and_pad 等通信原语。

一句话：

MF 当前实现已经给出了完整两段式 CP forward 的可运行参考：MF 选择 module 和输入切分边界，hook 注入 CP meta，attention construct 委托到一个专用 CP forward，CP forward 内部封装所有通信和全局 top-k。HP 迁移时应该保留这个职责边界，把实现主体从 MF local-Tensor small-op 升级为 HP 管理 meta、fused op 和 DTensor/DistributedOp 的版本。

yide12Collaborators

Jun 5 Close issue

yide12Collaborators

Jun 5 Comments：

MindFormers 两段式 DeepSeek V4 CP 实现拆解

本文只关注“完整两段式 CP 能力”这一类问题：

MF 如何选择 attention 边界、注入 CP meta，并在 CP forward 内封装 window KV halo、ratio4 compressor halo、compressed KV all-gather、select-and-pad、global top-k、KL loss 等细节。

不展开 PR717 custom op 的 DistributedOp/layout 语义。

参考实现：

/home/wyd/code_review_deepseek_v4_cp/mindformers_pr8311_latest
PR8311 head: a74de4127e8e72ae633ca21c3605e4a956d34399
主要文件：
- mindformers/pynative/base_models/gpt/parallelize.py
- mindformers/pynative/distributed/context_parallel.py
- mindformers/pynative/transformers/experimental_attention_variant/csa.py
- mindformers/pynative/transformers/experimental_attention_variant/csa_context_parallel.py

1. 先给结论

MF 当前没有调用一个来自 HP 的 DeepSeekV4CompressedSparseAttentionCPForward 对象。它是在 MF 仓内实现了同构能力：

MF parallelize 选择 CompressedSparseAttention module
  -> 挂一个 light forward-pre-hook
  -> hook 只注入 _cp_ctx = (cp_size, cp_rank, cp_group)
  -> CompressedSparseAttention.construct 发现 _cp_ctx 后委托给 two_stage_cp_naive_attention
  -> two_stage_cp_naive_attention 封装完整 Stage-1 + Stage-2

flowchart TB
    A["parallelize_gptmodel"] --> B["Phase 3: apply_context_parallel_attention"]
    B --> C{"core_attention 是 CompressedSparseAttention?"}
    C -->|否| D["普通 attention<br/>套已有 ContextParallelAttentionStyle"]
    C -->|是| E["CompressedSparseAttentionContextParallel(cp_mesh)._apply"]
    E --> F["注册 forward-pre-hook"]
    F --> G["hook 注入 _cp_ctx"]
    G --> H["CompressedSparseAttention.construct"]
    H --> I{"_cp_ctx and cp_size > 1?"}
    I -->|否| J["原单卡路径"]
    I -->|是| K["two_stage_cp_naive_attention"]
    K --> L["Stage-1: global compressed KV"]
    L --> M["Stage-2: local Q attention"]

2. MF 如何决定替换哪个 forward

入口在 parallelize_gptmodel：

Phase 1 TP
Phase 2 EP
Phase 3 CP attention hook
Phase 4 recompute/swap
Phase 5 FSDP
Phase 6 CP root I/O hook

CP attention hook 的核心逻辑在 apply_context_parallel_attention：

遍历 gpt_model.decoder.layers。
在每个 layer 内查找 self_attention.*core_attention。
如果该 module 是 CompressedSparseAttention，安装 DeepSeek V4 专用 two-stage CP hook。
否则普通 attention 继续走原有 CP style。

这点和目标 HP 接口边界一致：

MF 自己遍历 MF 模型。
MF 自己知道哪个 module 是 DeepSeek V4 CSA/HCA attention。
HP 不应该扫描 MF 模型，也不应该按 MF 私有命名猜 hook 位置。

当前 MF 代码等价于未来这样的调用：

for layer in decoder.layers:
    csa = find_core_attention(layer)
    if isinstance(csa, CompressedSparseAttention):
        hp.patch_deepseek_v4_compressed_sparse_attention_forward(
            csa,
            cp_mesh=cp_mesh,
            config=config,
        )

区别只是：现在 patch 逻辑由 MF 自己的 CompressedSparseAttentionContextParallel 完成。

3. MF 如何传 CP meta

CompressedSparseAttentionContextParallel 初始化时从 cp_mesh 取三项：

cp_size = cp_mesh.size()
cp_rank = cp_mesh.get_local_rank()
cp_group = cp_mesh.get_group()

然后 _apply(module) 注册 pre-hook：

new_kwargs["_cp_ctx"] = (cp_size, cp_rank, cp_group)

这个 hook 有两个重要约束：

hook 不做 compute。
hook 不做 communication。

原因是通信必须发生在 construct 的 autograd 上下文里，否则容易遇到 hook 内 recompute/top-cell 相关问题，也不好保证反向 collective 对称。

sequenceDiagram
    participant P as MF parallelize
    participant M as CompressedSparseAttention
    participant H as CP pre-hook
    participant C as construct
    participant F as two_stage_cp_naive_attention

    P->>M: register_forward_pre_hook
    M->>H: forward call enters hook
    H->>H: build kwargs['_cp_ctx']
    H-->>C: no compute, no comm
    C->>C: check _cp_ctx
    C->>F: delegate heavy CP forward

建议 HP 目标 meta 至少包含：

cp_size
cp_rank
cp_group
global_seq_len
local_seq_start
local_seq_len
position_ids_local
cu_seqlens_padded
cu_seqlens_real
cu_seqlens_cmp_kv

4. MF 如何保证输入已经是 local sequence

两段式 attention forward 不负责切 input_ids / labels / loss_mask / position_ids。MF 在 root model 边界另挂了 ContextParallelModelIOStyle。

它处理的输入名：

input_ids
position_ids
attention_mask
decoder_input
labels
loss_mask
actual_seq_len

普通 BS 路径：

[B, S] -> [B, S / CP]

TND/packed 路径检测到 actual_seq_len 后：

[B, S] -> reshape [1, B*S] -> [1, B*S / CP]

当前实现要求 position_ids 必须存在。main Q/K RoPE 依赖切分后的 local position_ids 保持全局位置语义。

flowchart TB
    A["全局 batch"] --> B["ContextParallelModelIOStyle"]
    B --> C["prepare_context_parallel_input"]
    C --> D["input_ids_local"]
    C --> E["labels_local"]
    C --> F["loss_mask_local"]
    C --> G["position_ids_local"]
    G --> H["GPTModel rotary_pos_emb(position_ids=...)"]
    H --> I["main Q/K 已按全局位置 RoPE"]
    I --> J["进入 CompressedSparseAttention 的 query/key/x/qr 已是 local shard"]

这也是未来 HP 接口要遵守的边界：HP 的 CP attention forward 可以校验输入 shape/meta，但不应该在 CompressedSparseAttention 内临时切全局 batch。

5. CP forward 的实际入口

CompressedSparseAttention.construct 增加 _cp_ctx=None 参数。

逻辑是：

if _cp_ctx is not None and cp_size > 1:
    if not force_unfused:
        raise ValueError
    return two_stage_cp_naive_attention(self, query, ori_kv, x, qr, *cp_ctx)
else:
    run original single-card construct

这等价于前向替换对象的 __call__：

def __call__(query, key, value, attention_mask, x=None, qr=None):
    if cp_disabled:
        return origin_construct(...)
    return construct_cp_two_stage(...)

当前限制：

CP>1 只支持 force_unfused_dsa=True。
fused npu_sparse_attn_sharedkv CP 主路径未实现。
因此 MF 当前是在 small-op naive 路径上验证算法，不是最终性能路径。

6. Stage-1：compressed KV 全局化

two_stage_cp_naive_attention 内部首先计算：

ratio = cell.compress_ratio
s_local = query.shape[0]
start = cp_rank * s_local
overlap = ratio == 4

6.1 ratio4 CSA：compressor halo

ratio4 是重叠压缩。边界 block 需要左邻最后 ratio 个 token。

当前实现没有用 P2P，而是：

x_bnd = all_gather(x[-ratio:])
halo_x = x_bnd[(cp_rank - 1) * ratio : cp_rank * ratio]

然后对 main compressor 和 indexer compressor 都执行：

local_slots = stage1_compress_with_halo(compressor, x, halo_x, cp_rank, overlap=True)
gathered = all_gather(local_slots)
gather_index = build_select_pad_gather_index(...)
cmp_global = select_and_pad_smallop(gathered, gather_index)

flowchart TB
    A["x local shard"] --> B["all_gather 每 rank x[-ratio:]"]
    B --> C["取左邻 halo_x"]
    C --> D1["cell.compressor<br/>stage1_compress_with_halo"]
    C --> D2["cell.indexer.compressor<br/>stage1_compress_with_halo"]
    D1 --> E1["all_gather local_slots"]
    D2 --> E2["all_gather local_slots"]
    E1 --> F1["select-and-pad"]
    E2 --> F2["select-and-pad"]
    F1 --> G1["cmp_main_global"]
    F2 --> G2["cmp_idx_global"]

slot 语义：

rank0：slot0 = PAD，后面是本 rank valid blocks。
rank>0：slot0 = boundary DUP，后面是本 rank valid blocks。
select-and-pad 删除 PAD/DUP/无效半组/padding block。

6.2 ratio128 HCA：非重叠 compressed KV all-gather

ratio128 不需要 halo_x：

local_slots = compressor(local_x, rope_pos_offset=rank * s)
gathered = all_gather(local_slots)
cmp_main_global = select_and_pad_smallop(gathered, gather_index)

HCA 没有 indexer，所以没有 cmp_idx_global。

6.3 当前 packed meta 缺口

build_select_pad_gather_index 支持：

cu_seqlens_padded
cu_seqlens_real
cp
ratio

但 runtime 当前调用是：

build_select_pad_gather_index(
    [0, s_local * cp_size],
    [0, s_local * cp_size],
    cp_size,
    ratio,
    overlap=overlap,
)

也就是只按单段 full-valid sequence 构造 gather index。packed/TND 的真实 cu_seqlens 语义有 helper 和 UT，但没有接进 two-stage CP runtime 主路径。

对 HP 来说，这是必须补齐的 CP metadata 管理点。

7. Stage-2：local Q attention

Stage-2 的原则：

Q 不 all-gather
输出不 scatter-back
每 rank 只算本地 query shard 的 output

7.1 window KV halo

对 window ori_kv：

hlen = min(window - 1, s_local)
ori_bnd = all_gather(ori_kv[-hlen:])
ori_ext, halo_len = assemble_window_halo(ori_bnd, ori_kv, cp_rank, hlen, window)

assemble_window_halo 支持 window - 1 > s_local，也就是 halo 可能跨多个左侧 rank，不只取直接左邻。

7.2 调用 `_construct_naive`

Stage-2 最后委托回 MF 已验证的 naive attention：

out = cell._construct_naive(
    query,
    ori_ext,
    x,
    qr,
    pos_offset=start,
    cmp_kv_global=cmp_main,
    ori_halo_len=halo_len,
    cmp_idx_global=cmp_idx,
)

_construct_naive 因 CP 新增的参数含义：

pos_offset: 本 rank local query 的全局起点。
cmp_kv_global: Stage-1 得到的全局 main compressed KV。
ori_halo_len: ori_ext 前面拼了多少 window halo。
cmp_idx_global: Stage-1 得到的全局 indexer compressed K。

flowchart TB
    A["local query"] --> H["cell._construct_naive"]
    B["ori_ext = left halo + local ori_kv"] --> H
    C["cmp_main_global"] --> H
    D["cmp_idx_global, ratio4 only"] --> H
    E["pos_offset = cp_rank * s_local"] --> H
    F["ori_halo_len"] --> H
    H --> I["local output shard"]

7.3 HCA global visibility

ratio128 没有 indexer。_construct_naive 调：

get_compress_topk_idxs(
    ratio,
    batch,
    local_seq_len,
    offset=len(ori_ext),
    pos_offset=global_start,
    n_compressed=len(cmp_main_global),
)

这样每个本地 query 的 compressed visibility 按全局位置计算。

7.4 CSA global top-k

ratio4 有 indexer。CP 下不能用本地 compressed K 做 top-k，必须用 cmp_idx_global：

q_idx, _, w_idx = indexer.forward_before_topk(
    x_local,
    qr_local,
    query_pos_offset=global_start,
)

topk = fused_qk_topk_naive(
    q_idx,
    cmp_idx_global,
    w_idx,
    eff_topk,
    global_causal_mask,
)

关键修正：

indexer Q 的 RoPE 用 query_pos_offset。
key 是 global indexer compressed K。
causal mask 用 global query positions。
top-k 宽度按 key length 截断，而不是 local query length。

8. KL loss 当前是怎么做的

ratio4 training 时，_construct_naive 会计算 naive KL：

compute_dsa_indexer_loss(...)
output = DSAIndexerLossAutoScaler(output, indexer_loss)
save_to_indexer_losses_tracker(indexer_loss, ...)

需要明确：MF 当前 two-stage CP forward 内没有显式做 CP all-reduce KL loss 标量。

源码注释给出的边界是：

naive compute_dsa_indexer_loss 不做内部 all-reduce。
tracker 是 single-card tracker。
cross-rank reduce 由 trainer/callback step end 负责。

因此你列的“KL loss all-reduce”更适合作为 HP 目标设计要求，而不是 MF 当前 CP forward 已实现事实。HP 迁移时应把 KL loss 分成两件事：

梯度路径：indexer loss 必须正确接入 output graph。
标量统计/scale：跨 CP rank 的 reduce、tracker 记录、loss scale 需要定义清楚。

flowchart TB
    A["ratio4 global top-k"] --> B["compute_dsa_indexer_loss"]
    B --> C["DSAIndexerLossAutoScaler<br/>接梯度"]
    B --> D["save_to_indexer_losses_tracker<br/>记录统计"]
    D --> E["trainer/callback<br/>外层 parallel loss 修正"]
    B -.当前 CP forward 内没有.-> F["显式 CP all-reduce KL"]

9. 反向通信对称性

MF 当前通信用函数式 all_gather_into_tensor。代码注释说明它的 backward 是 reduce-scatter VJP。

为了避免某些 gathered tensor 在某些 rank 上没有进入 loss graph，导致 backward collective 不对称，MF 做了 zero-weight tie：

for t in ties:
    out = out + ops.cast(t.sum(), out.dtype) * 0.0

被 tie 的 tensor 包括：

ratio4 x_bnd
window ori_bnd

这不是数值修正，而是 collective 对称性修正。HP 实现时必须保留同等语义；如果未来把 halo all-gather 换成 P2P，需要单独实现 autograd wrapper。

10. MF 测试如何证明 CP forward 正确

MF 的测试分三层。

flowchart TB
    A["helper / unit tests"] --> A1["gather_index END-based"]
    A --> A2["cmp_cu_seqlens"]
    A --> A3["select-and-pad"]

    B["single-process CP reconstruction"] --> B1["Stage-1 cmp == CP=1"]
    B --> B2["RoPE offset == CP=1"]
    B --> B3["ratio0 / ratio128 / ratio4 output == CP=1"]

    C["real multi-card ST"] --> C1["2-card Stage-1 real all-gather"]
    C --> C2["whole-net CP=2 smoke"]
    C --> C3["same checkpoint CP=1 vs CP=2/4 loss equivalence"]

    A --> B --> C

最关键的测试：

test_csa_context_parallel.py
- 验证 Stage-1、window halo、HCA、CSA、RoPE offset、top-k cap。
test_two_card_cp_comm.py
- 真实 2 卡通信验证 global compressed KV 对齐 CP=1。
test_cp_loss_equivalence.py
- 同一 checkpoint 下 CP=1 vs CP=2/4 逐 step loss 对齐。

这正对应 HP 未来需要提供的“CP=1 与 CP>1 数值对齐测试”。

11. 和你列出的 HP 第二类能力逐项对照

目标能力	MF 当前实现方式	HP 化建议
压缩注意力前向替换接口	`CompressedSparseAttentionContextParallel` + `_cp_ctx` + `two_stage_cp_naive_attention`	HP 提供 `DeepSeekV4CompressedSparseAttentionCPForward` 或 patch 函数，MF 负责选择 module
HP 不遍历 MF 模型	当前 MF 自己遍历 `decoder.layers` 并识别 `CompressedSparseAttention`	保持该边界，HP 只处理显式传入的 CSA module
MF 决定 sequence shard 边界	`ContextParallelModelIOStyle` 挂在 root model I/O	HP CP forward 不应临时切全局 batch，只消费 local shard 和 meta
CP metadata	当前只有 `_cp_ctx=(size, rank, group)`，local range 用 `rank * s_local` 推导	HP 需要扩展为显式 meta，特别是 packed `cu_seqlens`
window KV halo	`ori_kv[-hlen:]` all-gather + `assemble_window_halo`	HP 可先复用 all-gather 语义，后续优化 P2P
ratio4 compressor halo	`x[-ratio:]` all-gather，取左邻 halo	HP 需要支持 main compressor 和 indexer compressor 两套 Stage-1
compressed KV all-gather	`_ag_seq(local_slots, cp_group)`	HP fused/DTensor 版本仍要保证 backward reduce-scatter 对称
select-and-pad	host `gather_index` + `cat zero row + index_select`	HP 可先 small-op，后续 fused select-and-pad
global top-k	local indexer Q/weights + global `cmp_idx` + global causal mask	HP 不能在本地 compressed K 上做 top-k
KL loss all-reduce	当前 CP forward 内没有显式 all-reduce，依赖外层统计/scale	HP 需要明确 Op-3/naive loss 的 reduce 与 tracker 语义
CP=1/CP>1 对齐测试	已有 helper、single-process golden、2-card comm、whole-net loss equivalence	HP 应保留同样分层测试

12. 对 HP 接口的直接启示

未来 HP 接口建议长这样：

def patch_deepseek_v4_compressed_sparse_attention_forward(
    csa_module,
    *,
    cp_mesh,
    cp_meta_provider,
):
    ...

其中：

csa_module 由 MF 传入。
cp_mesh 提供 rank/size/group。
cp_meta_provider 提供 local/global sequence、position ids、packed cu_seqlens。
HP patch 后的 forward 内部完整封装 Stage-1 + Stage-2。
MF 不直接调用 exchange_halo / all_gather_cmp / select_and_pad 等通信原语。

一句话：

MF 当前实现已经给出了完整两段式 CP forward 的可运行参考：MF 选择 module 和输入切分边界，hook 注入 CP meta，attention construct 委托到一个专用 CP forward，CP forward 内部封装所有通信和全局 top-k。HP 迁移时应该保留这个职责边界，把实现主体从 MF local-Tensor small-op 升级为 HP 管理 meta、fused op 和 DTensor/DistributedOp 的版本。

yide12Collaborators

Jun 11 Reopen issue

yide12Collaborators

Jun 11 Comments：

HP 适配 DeepSeek V4 CP 的具体实现设计

本文是 /home/wyd/code_review_deepseek_v4_cp/设计文档.md 的工程细化版。目标不是重新解释职责边界，而是把 MF PR8311 里的两阶段 CP 语义，落成 HP 侧可以直接开工的模块、接口、数据结构、通信流程和测试计划。

参考代码：

MF PR8311：/home/wyd/code_review_deepseek_v4_cp/mindformers_pr8311_latest
HP 开发分支：/home/wyd/code_review_deepseek_v4_cp/hp_feature
MF 走读笔记：/home/wyd/code_review_deepseek_v4_cp/tmp.py

1. 总体结论

HP 需要新增的是一个 DeepSeek V4 CSA 两阶段 CP forward 能力，而不是一堆让 MF 自己拼的通信小函数。

MF 调用边界应该是：

hp.patch_deepseek_v4_compressed_sparse_attention_forward(
    csa_module,
    cp_mesh=cp_mesh,
    config=cp_config,
)

MF 负责选择哪些 CompressedSparseAttention module 需要 patch；HP 负责 patch 后完整接管 CP>1 的 attention forward。

核心 forward：

local Q / local ori_kv / local x / local qr
        |
        |-- window ori_kv halo
        |-- Stage-1: global compressed KV
        |-- ratio=4: global indexer top-k
        |
        v
local Q 在 [window halo + global compressed KV] 上算 attention
        |
        v
local output shard

最重要的工程判断：

已有 DSA ContextParallel style / DistributedOp 可以复用设计模式，
但不能直接等价替代 DeepSeek V4 两阶段 CP forward。

原因：

现有 DSAIndexerContextParallel / DSASparseAttentionContextParallel 是 op boundary hook，语义是“query shard，key/value gather”。
DeepSeek V4 CP 的关键不是普通 KV gather，而是：
- window KV halo；
- compressor halo；
- compressed KV all-gather；
- select-and-pad 去掉 PAD/DUP/无效 compressed slot；
- ratio=4 在 global compressed K 上做 top-k；
- sharedKV sparse attention 使用 prepared window KV + global compressed KV。

因此 HP 要做的是：

DeepSeekV4 CP forward 编排
    +
PR717 三个 custom op 的 DistributedOp/layout 语义

2. 当前 HP 分支能力盘点

当前 hp_feature 分支里已经有几类可复用基础。

2.1 可复用：ContextParallel 代码风格

已有文件：

hyper_parallel/core/context_parallel/context_parallel.py
hyper_parallel/core/context_parallel/dsa_context_parallel.py
hyper_parallel/core/context_parallel/async_dsa_context_parallel.py

可复用点：

ParallelStyle.apply(module, device_mesh) 的接口风格。
forward pre-hook / forward hook 的 patch 机制。
_ensure_1d(device_mesh)、DeviceMesh.get_group() 这类 CP mesh 获取方式。
_to_cp_dtensor(...)、Shard(seq_dim)、Replicate() 的 DTensor 标注方式。
async DSA 里 producer/consumer all-gather slot 的结构，后续可借鉴做异步 compressed KV gather。

不可直接复用点：

DSASparseAttentionContextParallel 会把 key/value 做 sequence replicate，这不等价于 DeepSeek V4 的 window halo。
DSAIndexerContextParallel 会把 key-side gather 成全局 K，但没有处理 compressor halo、select-and-pad、packed compressed cu_seqlens。
这些 style 不知道 MF CompressedSparseAttention 内部的 ratio/window/compressor/indexer 语义。

2.2 可复用：DistributedOp 写法

已有文件：

hyper_parallel/core/shard/ops/parallel_lightning_indexer.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_flash_attention.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_lightning_indexer_grad_kl_loss.py
hyper_parallel/core/shard/ops/yaml/lightning_indexer_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_flash_attention_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_lightning_indexer_grad_kl_loss_ops.yaml

可复用点：

preprocess -> infer_layout -> get_expand_impl 三段式。
对 BSND/TND layout 的参数归一化。
query 侧 Shard(seq)、key 侧 Replicate()、loss 输出 Partial(sum) 的布局推导。
TND packed seq lens 的调整模式。

缺口：

当前分支尚未看到 PR717 目标名字对应的 wrapper 和 DistributedOp：

npu_lightning_indexer_v2
npu_sparse_attn_sharedkv
npu_sparse_lightning_indexer_grad_kl_loss_v2

因此 DeepSeek V4 CP 需要新增这三个 op 的 HP wrapper / platform custom op / DistributedOp 规则。

2.3 可复用：通信接口

已有平台接口：

platform.differentiable_all_gather_concat(...)
platform.differentiable_reduce_scatter(...)
platform.all_gather_into_tensor(...)
platform.reduce_scatter_tensor(...)
platform.all_gather_single(...)
platform.reduce_scatter_single(...)

首版建议：

window halo:       all-gather boundary tail，然后本地 slice
compressor halo:   all-gather ratio 个边界 token，然后本地 slice
compressed KV:     all-gather local compressed slots
KL loss:           all-reduce / Partial(sum) reduce

P2P 可以作为后续优化，不进入首版正确性闭环。

3. 推荐新增文件

建议优先集中在 hyper_parallel/core/context_parallel/ 下，避免一开始拆太碎。

hyper_parallel/core/context_parallel/deepseek_v4_context_parallel.py

首版可以一个文件放完 public API、config、meta、forward class 和 small-op helper。等跑通后再拆：

hyper_parallel/core/context_parallel/deepseek_v4_cp_meta.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_indices.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_collectives.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_forward.py

custom op 层新增：

hyper_parallel/custom_ops/experimental/experimental_ops.py
hyper_parallel/platform/mindspore/custom_ops/custom_ops.py
hyper_parallel/platform/mindspore/custom_ops/custom_op_impl.py
hyper_parallel/platform/mindspore/custom_ops/*.cc
hyper_parallel/core/shard/ops/parallel_npu_lightning_indexer_v2.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_attn_sharedkv.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py
hyper_parallel/core/shard/ops/yaml/npu_lightning_indexer_v2_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_attn_sharedkv_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_lightning_indexer_grad_kl_loss_v2_ops.yaml

__init__ 导出：

hyper_parallel/core/context_parallel/__init__.py
hyper_parallel/__init__.py

新增 public symbol：

DeepSeekV4CPConfig
DeepSeekV4CPContext
DeepSeekV4CompressedSparseAttentionCPForward
patch_deepseek_v4_compressed_sparse_attention_forward

4. Public API 设计

4.1 Patch API

def patch_deepseek_v4_compressed_sparse_attention_forward(
        csa_module,
        *,
        cp_mesh,
        config,
):
    ...

含义：

csa_module：MF 选中的 CompressedSparseAttention Cell。
cp_mesh：HP DeviceMesh，首版只支持 1-D CP mesh。
config：HP DeepSeek V4 CP 配置。

原 csa_module 本体

不替换 Cell 对象，只 patch construct，保留 MF module 的参数注册和 state_dict。

4.2 Config

@dataclass
class DeepSeekV4CPConfig:
    enabled: bool = True
    mode: str = "colossal"
    input_layout: str = "SBNHD"       # MF CSA 边界: [S, B, N, D]
    op_layout: str = "BSND"           # NPU op 边界
    seq_dim: int = 0                  # MF CSA 边界 sequence dim
    op_seq_dim: int = 1               # BSND op sequence dim
    window_size: int | None = None    # None 表示从 csa_module 读取
    use_fused_ops: bool = True
    use_allgather_halo: bool = True
    strict: bool = True
    debug_force_smallop: bool = False

首版只支持：

mode == "colossal"
input_layout == "SBNHD"
op_layout == "BSND"

Ulysses/head shard 不进入 DeepSeek V4 CP 首版，因为 DSv4 custom ops 要求 head/head_dim/topk 维保持 replicated。

4.3 Runtime Context

DeepSeekV4CPContext 是 HP 在 forward 内部推导出来的运行时上下文：

@dataclass
class DeepSeekV4CPContext:
    cp_mesh: DeviceMesh
    cp_group: object
    cp_rank: int
    cp_size: int

    global_seq_len: int
    local_seq_start: int
    local_seq_end: int
    local_seq_len: int

    batch_size: int
    seq_dim: int
    op_layout: str

    compress_ratio: int
    window_size: int

    # packed sequence，可空
    cu_seqlens_q: object | None = None
    cu_seqlens_kv: object | None = None
    cu_seqlens_cmp: object | None = None
    cu_seqlens_padded: object | None = None
    cu_seqlens_real: object | None = None

DeepSeekV4CPContext 应由 HP 构建，不应要求 MF 直接填完整对象。

非 packed 等长序列只需要：

cp_mesh
query.shape[0]

HP 可推导：

cp_rank
cp_size
local_seq_len = query.shape[0]
global_seq_len = local_seq_len * cp_size
local_seq_start = cp_rank * local_seq_len
local_seq_end = local_seq_start + local_seq_len

packed/TND 场景需要额外知道文档边界：

actual_seq_len / cu_seqlens

这份信息不应通过 side-channel callback 传入 HP。它应该是正常 forward 数据流的一部分：

方案 A：MF 将 actual_seq_len/cu_seqlens 正常透传到 CompressedSparseAttention CP forward。
方案 B：如果 CSA 边界拿不到 actual_seq_len，则 MF 选择更高一层 DeepSeekV4HybridAttention 作为 HP CP hook 边界。

有了 actual_seq_len 后，HP 可继续推导：

每个 rank 的 global token range
每个 rank 覆盖哪些 doc
window halo 是否跨 doc
ratio=4 compressor halo 是否跨 doc
doc 内 query position
compressed cu_seqlens
select-and-pad gather_index

4.4 MF 调用方式

MF 侧推荐：

def apply_deepseek_v4_context_parallel(model, cp_mesh, config):
    for csa_module in model.iter_deepseek_v4_csa_modules():
        hp.patch_deepseek_v4_compressed_sparse_attention_forward(
            csa_module,
            cp_mesh=cp_mesh,
            config=config,
        )
    return model

HP 不提供：

hp.apply_deepseek_v4_context_parallel(model)

因为 HP 不应该遍历 MF 模型，也不应该根据 MF 结构猜 hook 位置。

5. Forward 替换对象

5.1 类结构

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __init__(self, parent, cp_mesh, config):
        self.parent = parent
        self.origin_construct = parent.construct
        self.cp_mesh = _ensure_1d(cp_mesh)
        self.config = config

    def __call__(self, query, key, value, attention_mask, x=None, qr=None, actual_seq_len=None):
        if not self.config.enabled or self._cp_size() == 1:
            return self.origin_construct(
                query, key, value, attention_mask, x=x, qr=qr
            )

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

Patch 函数：

def patch_deepseek_v4_compressed_sparse_attention_forward(...):
    cp_forward = DeepSeekV4CompressedSparseAttentionCPForward(...)
    csa_module._hp_origin_construct = csa_module.construct
    csa_module._hp_deepseek_v4_cp_forward = cp_forward

    def construct(query, key, value, attention_mask, x=None, qr=None):
        return cp_forward(query, key, value, attention_mask, x=x, qr=qr)

    csa_module.construct = construct
    return csa_module

5.2 入口输入 shape

MF CompressedSparseAttention.construct 边界是 seq-first：

query: [s_local, B, N, D]
key:   [s_local, B, 1, D] 或 [s_local, B, D]
value: [s_local, B, 1, D] 或 [s_local, B, D]
x:     [s_local, B, H]
qr:    [s_local, B, q_lora_rank]

HP 入口必须强校验：

query.shape[0] == cp_ctx.local_seq_len
key.shape[0]   == cp_ctx.local_seq_len
x.shape[0]     == cp_ctx.local_seq_len
qr.shape[0]    == cp_ctx.local_seq_len

ratio 约束：

ratio=0:   x/qr 可不参与 compressed 分支
ratio=128: x 必须存在
ratio=4:   x 和 qr 必须存在，且 parent.indexer 必须存在

6. CP Forward 主流程

主流程伪代码：

def _construct_cp(self, query, key, value, attention_mask, x, qr, cp_ctx):
    parent = self.parent
    ratio = self._read_compress_ratio(parent)

    ori_kv = self._extract_shared_kv(key, value)

    ori_ext, ori_halo_len = self._prepare_window_kv_halo(
        ori_kv, cp_ctx, window_size=parent.window_size
    )

    cmp_main = None
    cmp_idx = None
    if ratio and ratio > 1:
        cmp_main = self._stage1_global_compressed_kv(
            parent.compressor, x, cp_ctx, purpose="main"
        )

    if ratio == 4 and parent.indexer is not None:
        cmp_idx = self._stage1_global_compressed_kv(
            parent.indexer.compressor, x, cp_ctx, purpose="indexer"
        )

    if self.config.debug_force_smallop:
        return self._stage2_smallop(
            query, ori_ext, x, qr,
            cmp_main, cmp_idx,
            ori_halo_len, cp_ctx,
        )

    return self._stage2_fused_sharedkv(
        query, ori_ext, x, qr,
        cmp_main, cmp_idx,
        ori_halo_len, cp_ctx,
    )

建议首个可验证版本先实现：

debug_force_smallop=True

即 HP 负责 Stage-1 和 window halo，Stage-2 先复用 MF module 的 _construct_naive(...) 或等价 small-op 实现，快速验证 CP=1/CP>1 数值对齐。之后再切到 PR717 fused sharedKV op。

7. Window KV Halo

7.1 功能

window halo 用来解决 rank 边界处 query 看不到左侧 KV 的问题。

例子：

CP = 2
s_local = 4
window_size = 4

rank0: 0 1 2 3
rank1: 4 5 6 7

rank1 的 token 4 要看：

1 2 3 4

所以 rank1 需要从 rank0 拿：

1 2 3

这段就是 window KV halo。

7.2 首版通信

首版用 all-gather boundary tail：

hlen = min(window_size - 1, s_local)
local_tail = ori_kv[-hlen:]                 # [hlen, B, D]
all_tails = all_gather_seq(local_tail)      # [cp_size * hlen, B, D]
left_halo = slice_left_context(all_tails, cp_rank, hlen, window_size)
ori_ext = cat([left_halo, ori_kv], dim=0)

shape：

ori_kv:    [s_local, B, D]
local_tail:[hlen, B, D]
all_tails: [cp_size * hlen, B, D]
left_halo: [halo_len, B, D]
ori_ext:   [halo_len + s_local, B, D]

hlen 是每个 rank 贡献多少 boundary，不是当前 rank 最终拿多少 halo：

hlen = min(window_size - 1, s_local)

当 window_size - 1 > s_local 时，当前 rank 的 left_halo 可能来自多个前序 rank；all-gather 后本地 slice 可以统一处理。

7.3 函数建议

def _prepare_window_kv_halo(ori_kv, cp_ctx, window_size):
    hlen = min(window_size - 1, cp_ctx.local_seq_len)
    if hlen <= 0 or cp_ctx.cp_rank == 0:
        ...
    all_tails = _all_gather_seq(ori_kv[-hlen:], cp_ctx)
    left_halo = _assemble_left_halo_from_rank_tails(
        all_tails,
        cp_rank=cp_ctx.cp_rank,
        hlen=hlen,
        need=window_size - 1,
    )
    return cat([left_halo, ori_kv], dim=0), left_halo.shape[0]

注意 packed sequence：

halo 不能跨文档。

首版非 packed 可先按连续全局 token slice；packed 版本必须按 cu_seqlens_real 对每个 doc 单独组 halo。

8. Stage-1：Global Compressed KV

Stage-1 的唯一目标：

让每个 rank 都拿到 CP=1 语义下的 global compressed KV。

Q 不在 Stage-1 压缩。x 分两路：

x -> attention projection -> Q / ori_kv
x -> compressor           -> compressed KV

8.1 ratio=0

无 compressed KV：

cmp_main = None
cmp_idx = None

只走 window halo + sharedKV attention。

8.2 ratio=128

非重叠压缩。

流程：

local_slots = compressor(
    x,
    rope_pos_offset=cp_ctx.local_compressed_start * ratio,
)
gathered = all_gather_seq(local_slots)
cmp_global = select_and_pad(gathered, gather_index)

非 packed 且 s_local % 128 == 0 时：

local_slots: [s_local / 128, B, D]
gathered:    [S / 128, B, D]
cmp_global:  [S / 128, B, D]

需要注意因果可见性：

query position p 只能看 compressed slot c < (p + 1) // ratio

例如 p=300，ratio=128：

visible slots: 0, 1
slot 2 压缩 256..383，包含未来 301..383，不能看。

8.3 ratio=4

重叠压缩，需要 compressor halo。

MF 当前语义可以抽象为：

1. 当前 rank 拿左侧 ratio 个 hidden states 作为 compressor halo
2. compressor 在 [halo_x, local_x] 上压缩
3. 每个 rank 生成等长 local slots
4. rank0 的第一个 slot 是 PAD
5. rank>0 的第一个 slot 是跨边界 DUP
6. all-gather 后用 select-and-pad 去掉 PAD/DUP/无效 slot

首版实现建议继续使用 all-gather boundary：

compress_hlen = ratio
x_tail = x[-compress_hlen:]
all_x_tails = all_gather_seq(x_tail)
halo_x = slice_left_context(all_x_tails, cp_rank, compress_hlen, need=ratio)
local_slots = stage1_compress_with_halo(compressor, x, halo_x, cp_ctx, overlap=True)
gathered = all_gather_seq(local_slots)
cmp_global = select_and_pad(gathered, gather_index)

main compressed KV 和 indexer compressed K 都要做：

cmp_main = stage1_global(parent.compressor)
cmp_idx  = stage1_global(parent.indexer.compressor)

shape：

x:             [s_local, B, H]
halo_x:        [ratio, B, H] 或更短
local_slots:   [local_slot_count_with_pad_dup, B, D]
gathered:      [cp_size * local_slot_count_with_pad_dup, B, D]
cmp_global:    [C, B, D]

cmp_idx:       [C, B, index_head_dim]

8.4 select-and-pad

ratio=4 的 overlap 会产生 PAD/DUP；packed sequence 还会产生文档 padding 无效块。因此 all-gather 后不能直接使用 gathered。

需要一个 index：

gather_index: [C]

做：

padded = cat([gathered, zero_sentinel], dim=0)
cmp_global = index_select(padded, dim=0, index=gather_index)

其中：

valid slot -> 指向 gathered 的真实位置
invalid    -> 指向 zero_sentinel

函数建议：

def build_compressed_gather_index(cp_ctx, ratio, overlap):
    ...

def select_and_pad_compressed(gathered, gather_index):
    ...

缓存建议：

key = (cp_size, global_seq_len, local_seq_len, ratio, packed_signature)

首版非 packed 可只支持：

S % (cp_size * ratio) == 0

packed 支持后再加入 cu_seqlens_real/padded。

9. Stage-2：Local Q Attention

Stage-2 输入：

query:     [s_local, B, N, D]
ori_ext:   [halo_len + s_local, B, D]
cmp_main:  [C, B, D] 或 None
cmp_idx:   [C, B, index_head_dim]，ratio=4 only

Stage-2 输出：

output: [s_local, B, N * D]

9.1 Small-op fallback 路径

用于首版对齐 MF PR8311。

如果 MF module 有 _construct_naive：

out = parent._construct_naive(
    query,
    ori_ext,
    x,
    qr,
    pos_offset=cp_ctx.local_seq_start,
    cmp_kv_global=cmp_main,
    ori_halo_len=halo_len,
    cmp_idx_global=cmp_idx,
)

语义：

kv_full = [ori_ext, cmp_main]
window_idxs 指向 kv_full 前半段
compressed_idxs 指向 kv_full 后半段
topk_idxs = [window_idxs, compressed_idxs]
unfused_compressed_sparse_attn(query, kv_full, topk_idxs)

这个 fallback 很适合做第一阶段 HP 验证：

HP 通信/metadata 正确性
MF CP=1 vs HP CP>1 数值对齐

但最终 fused 目标不应长期依赖 MF 私有 _construct_naive。

9.2 Fused sharedKV 路径

最终目标是调用 PR717 的：

npu_sparse_attn_sharedkv

输入准备：

q_local:        local Q
ori_kv_prepared: [left_halo, local_ori_kv]
cmp_kv_global: global compressed KV
cmp_sparse_indices:
  ratio=0:   None
  ratio=128: None，op 内部按 cmp_ratio/causal 处理 dense compressed visibility
  ratio=4:   global top-k indices from npu_lightning_indexer_v2

调用示意：

out = npu_sparse_attn_sharedkv(
    q_local_dt,
    ori_kv_prepared_dt,
    cmp_kv_global_dt,
    cmp_sparse_indices_dt,
    sinks=parent.attn_sink,
    scale_value=parent.softmax_scale,
    actual_seq_qlen=cp_ctx.cu_seqlens_q,
    actual_seq_klen=cp_ctx.cu_seqlens_kv,
    actual_seq_cmp_klen=cp_ctx.cu_seqlens_cmp,
    layout=cp_ctx.op_layout,
    cmp_ratio=ratio or None,
    ori_win_left=parent.ori_win_left,
    ori_win_right=parent.ori_win_right,
    local_seq_start=cp_ctx.local_seq_start,
    ori_halo_len=halo_len,
)

注意：

ori_kv_prepared 不是普通全局 replicated KV。
它是 HP 已经拼好的 [left_halo, local_ori_kv]。
DistributedOp 不能再把它当普通 key/value 做 all-gather。

因此 parallel_npu_sparse_attn_sharedkv.py 的布局规则要有 prepared window KV 语义。

10. ratio=4 Global Top-k

ratio=4 不能只在本 rank compressed K 上 top-k，必须使用 global indexer compressed K：

local q_idx / local weights
    +
global cmp_idx
    +
global causal mask / actual seq meta
    |
    v
global top-k compressed indices

10.1 q_idx / weights

复用 MF indexer 的前半段逻辑：

q_idx, _, w_idx = parent.indexer.forward_before_topk(
    stop_gradient(x),
    stop_gradient(qr),
    query_pos_offset=cp_ctx.local_seq_start,
)

shape：

q_idx: [s_local, B, index_n_heads, index_head_dim]
w_idx: [s_local, B, index_n_heads]

如果 fused op 需要 BSND：

q_idx -> [B, s_local, index_n_heads, index_head_dim]
w_idx -> [B, s_local, index_n_heads]
cmp_idx_global -> [B, C, index_n_heads or 1, index_head_dim]  # 以算子真实签名为准

10.2 调用 `npu_lightning_indexer_v2`

topk_indices, topk_scores = npu_lightning_indexer_v2(
    q_idx_dt,
    cmp_idx_global_dt,
    w_idx_dt,
    actual_seq_qlen=cp_ctx.cu_seqlens_q,
    actual_seq_klen=cp_ctx.cu_seqlens_cmp,
    layout=cp_ctx.op_layout,
    sparse_count=parent.indexer.index_topk,
    sparse_mode=3,
    cmp_ratio=4,
    query_pos_offset=cp_ctx.local_seq_start,
)

布局：

q_idx:          Shard(seq)
w_idx:          Shard(seq)
cmp_idx_global: Replicate()
topk_indices:   Shard(seq)
topk_scores:    Shard(seq)

输出 indices 是 compressed slot id，不是 kv_full 后半段 id。small-op fallback 里需要加 offset；fused sharedKV op 通常直接接收 compressed slot id。

10.3 Causal mask

因果规则：

compressed slot c 对 query global position p 可见 iff c < (p + 1) // ratio

packed sequence 下要按每个 doc 内位置算，不能用跨 doc 全局 token 直接除。

11. KL Loss

ratio=4 训练态需要 indexer KL loss。

目标 op：

npu_sparse_lightning_indexer_grad_kl_loss_v2

输入布局：

query:         Shard(seq)
key:           Replicate()  # main compressed/global key 或 sharedKV key，以 op 签名为准
query_index:   Shard(seq)
key_index:     Replicate()  # cmp_idx_global
weights:       Shard(seq)
sparse_indices:Shard(seq)
softmax_max:   Shard(seq)
softmax_sum:   Shard(seq)

输出：

d_query_index: Shard(seq)
d_key_index:   Replicate() 或 Partial(sum) 后 reduce
d_weights:     Shard(seq)
loss:          Partial(sum) -> all-reduce 后全局标量

HP 需要保证：

CP>1 的 KL loss 数值等价 CP=1。

工程上可沿用现有 DSAIndexerLossContextParallel 对 d_key_index 和 loss partial 的处理思路，但 op 名称、参数和 v2 语义要重新实现。

12. DistributedOp 规则

12.1 `parallel_npu_lightning_indexer_v2.py`

约束：

query_index: [B, S1, N1, D]，S1 可 Shard，N1/D replicated
key_index:   [B, S2, N2, D]，S2/N2/D replicated
weights:     [B, S1, N1]，B/S1 与 query_index 一致，N1 replicated

输出：

topk_indices: [B, S1, N2, sparse_count]
topk_scores:  [B, S1, N2, sparse_count]

布局：

B/S1 继承 query_index
N2/sparse_count replicated

CP expand：

key_index 已由 Stage-1 变成 global replicate，不再按普通 causal window 截 key。
只需要根据 local_seq_start / actual_seq_qlen / cmp_ratio 处理 causal visibility。

12.2 `parallel_npu_sparse_attn_sharedkv.py`

输入角色：

query:              Shard(seq)
ori_kv_prepared:    local prepared window KV with halo
cmp_kv_global:      Replicate()
cmp_sparse_indices: Shard(seq)，ratio=4 only

输出：

attention_out: Shard(seq)
softmax/lse:   Shard(seq) 或 op 定义布局

重点：

ori_kv_prepared 不能由 DistributedOp 自动 all-gather。
halo 已由 DeepSeekV4 CP forward 准备好。

建议在 wrapper 参数中显式标识：

ori_kv_is_prepared=True
ori_halo_len=halo_len
local_seq_start=cp_ctx.local_seq_start

如果用户直接拿普通 DTensor KV 调这个 op，应报错或要求走普通 DSA/SFA op。

12.3 `parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py`

复用现有 sparse KL loss DistributedOp 的大方向：

d_query_index: 继承 query_index layout
d_key_index:   key_index layout + CP Partial(sum)
d_weights:     继承 weights layout
loss:          scalar Partial(sum) over CP/DP

需要额外处理：

key_index 是 global compressed K，不是普通 full sequence K。
actual_seq_klen 应是 compressed cu_seqlens。

13. CP Meta 和 Packed Sequence

非 packed 首版：

global_seq_len = local_seq_len * cp_size
local_seq_start = cp_rank * local_seq_len
local_seq_end = local_seq_start + local_seq_len
cu_seqlens_q = None
cu_seqlens_cmp = None

packed 版本必须处理：

1. window halo 不跨文档
2. ratio=4 compressor halo 不跨文档
3. compressed cu_seqlens 按 doc 内 real length floor(real_len / ratio)
4. causal visible slot 按 doc 内 query position 计算
5. select-and-pad index 按 doc 生成

packed meta 的来源规则：

actual_seq_len / cu_seqlens 必须走正常 forward 数据流。

也就是说：

1. 如果 CompressedSparseAttention 边界能拿到 actual_seq_len，
   HP patch 这个边界。

2. 如果 CompressedSparseAttention 边界拿不到 actual_seq_len，
   MF 应选择更高一层 DeepSeekV4HybridAttention 作为 HP CP hook 边界，
   由这个边界把 actual_seq_len 和 query/key/x/qr 一起交给 HP。

HP CP forward 拿到 actual_seq_len/cu_seqlens 后，再自己推导：

local/global token range
doc 内 position
window halo 边界
compressor halo 边界
compressed cu_seqlens
select-and-pad gather_index

这里不引入 meta_provider、全局变量或 side-channel callback。完整
DeepSeekV4CPContext 是 HP 内部派生物，不由 MF 填。

14. 通信实现细节

14.1 `_all_gather_seq`

内部统一封装：

def _all_gather_seq(tensor, cp_ctx):
    return platform.differentiable_all_gather_concat(
        tensor,
        group=cp_ctx.cp_group,
        concat_size=cp_ctx.cp_size,
        concat_dim=0,
    )

MindSpore 平台也可以直接使用：

comm_func.all_gather_into_tensor(None, tensor.contiguous(), group=group)

但 HP 内部优先走 platform 抽象，方便 Torch/MindSpore 双后端和 autograd。

14.2 为什么首版不用 P2P

all-gather 逻辑：

每个 rank 发一样大小的 boundary
每个 rank 收所有 boundary
每个 rank 本地 slice

P2P 逻辑：

每个 rank 要算自己从哪些前序 rank 收、给哪些后续 rank 发、每个 peer 发多少、backward 怎么还梯度。

DeepSeek V4 CP 首版优先正确性，all-gather 更适合。

14.3 backward

首版依赖 differentiable all-gather 的 VJP：

forward all-gather
backward reduce-scatter / slice grad accumulate

select-and-pad 的 backward：

index_select forward
scatter_add backward

如果 MindSpore 小算子没有自动满足，需要实现自定义 grad 或使用已有 differentiable gather/scatter 组合。

15. 与 MF PR8311 的逐项映射

MF PR8311 能力	HP 目标实现位置
`two_stage_cp_naive_attention` 总控	`DeepSeekV4CompressedSparseAttentionCPForward._construct_cp`
window KV halo all-gather	`_prepare_window_kv_halo`
`_gather_compressed(cell.compressor)`	`_stage1_global_compressed_kv(parent.compressor)`
`_gather_compressed(cell.indexer.compressor)`	`_stage1_global_compressed_kv(parent.indexer.compressor)`
`stage1_compress_with_halo`	`_compress_local_slots_with_halo`
`build_select_pad_gather_index`	`build_compressed_gather_index`
`select_and_pad_smallop`	`select_and_pad_compressed`
`_construct_naive(... cmp_kv_global ...)`	small-op fallback / fused sharedKV path
ratio=4 `fused_qk_topk_naive`	`npu_lightning_indexer_v2`
`compute_dsa_indexer_loss`	`npu_sparse_lightning_indexer_grad_kl_loss_v2`

16. 开发里程碑

M1：HP small-op CP forward 跑通

范围：

新增 patch API。
新增 DeepSeekV4CPContext。
实现 window halo all-gather。
实现 ratio=128 global compressed KV。
实现 ratio=4 compressor halo + all-gather + select-and-pad。
Stage-2 先调用 MF _construct_naive 或 HP 等价 small-op。

验收：

ratio=0/4/128
CP=1 vs CP=2/4
forward output 对齐

M2：接入 PR717 v2/custom sharedKV ops

范围：

新增三个 experimental op wrapper。
MindSpore platform custom op 接入。
三个 DistributedOp + YAML 注册。
Fused Stage-2 使用 npu_sparse_attn_sharedkv。
ratio=4 top-k 使用 npu_lightning_indexer_v2。

验收：

small-op fallback vs fused path 对齐
CP=1 fused vs CP>1 fused 对齐

M3：KL loss 和反向

范围：

接入 npu_sparse_lightning_indexer_grad_kl_loss_v2。
loss all-reduce / Partial(sum)。
d_key_index CP partial reduce。
window halo / compressed gather / select-pad backward 验证。

验收：

loss 对齐
主要参数梯度对齐

M4：packed sequence

范围：

meta provider 支持 packed cu_seqlens。
window halo 不跨 doc。
compressor halo 不跨 doc。
compressed cu_seqlens per-doc。
ratio=4 select-and-pad per-doc。

验收：

多 doc packed batch
不同长度
rank 边界切在 doc 内 / doc 边界上 / padding 区

M5：异步和性能优化

范围：

window halo P2P 可选实现。
compressed KV all-gather async overlap。
indexer compressed K gather 与 q_idx 计算 overlap。
gather_index cache。

验收：

数值不变
通信时间下降
无额外峰值内存爆炸

17. 测试计划

17.1 纯函数 UT

新增：

tests/ut/core/context_parallel/test_deepseek_v4_cp_indices.py

覆盖：

hlen = min(window_size - 1, s_local)。
window halo 多 rank slice。
window_size - 1 < s_local 只需要前一个 rank。
window_size - 1 > s_local 需要多个前序 rank。
ratio=128 causal visible compressed slots。
ratio=4 select-and-pad index 去 PAD/DUP。

17.2 CP forward small-op UT/ST

新增：

tests/mindspore/st/context_parallel/deepseek_v4_cp_smallop.py
tests/mindspore/st/context_parallel/test_deepseek_v4_cp_smallop_msrun.py

覆盖：

ratio=0, window only
ratio=128, dense compressed
ratio=4, indexer top-k
CP=1/2/4
window 跨 0/1/多 rank

对齐标准：

CP=1 全量输入 reference
CP>1 本地输出 all-gather 后与 reference 对齐

17.3 DistributedOp UT

新增：

tests/ut/core/shard/test_parallel_npu_lightning_indexer_v2.py
tests/ut/core/shard/test_parallel_npu_sparse_attn_sharedkv.py
tests/ut/core/shard/test_parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py

覆盖：

layout inference。
非法 head/head_dim/topk sharding 报错。
query Shard(seq) 输出仍 Shard(seq)。
key/compressed K 必须 replicated。
KL loss 输出 partial reduce。

17.4 Fused ST

新增：

tests/mindspore/st/context_parallel/deepseek_v4_cp_fused.py
tests/mindspore/st/context_parallel/test_deepseek_v4_cp_fused_msrun.py

覆盖：

small-op fallback vs fused
CP=1 fused vs CP>1 fused
forward + backward

18. 风险和约束

18.1 不要把通信原语暴露给 MF

不推荐：

halo = hp.exchange_window_halo(...)
cmp = hp.all_gather_compressed_kv(...)
idx = hp.build_select_pad(...)

18.2 不要只加 DistributedOp 就结束

如果只给 npu_lightning_indexer_v2 / npu_sparse_attn_sharedkv 加 DistributedOp，而不改 CompressedSparseAttention.construct 的 CP forward，原 module 仍会先在本地 shard 上做 compressor/indexer：

local compressed KV 已经错了，
底层 op 再聪明也补不回 CP=1 语义。

18.3 prepared window KV 不是普通 replicated KV

DeepSeek V4 sharedKV op 的 ori_kv 输入在 CP forward 里已经变成：

[left_halo, local_ori_kv]

它不是：

all-gather 后的 full ori_kv

DistributedOp 规则必须区分这个语义。

18.4 packed sequence 是单独难点

非 packed 跑通不代表 packed 正确。

packed 下最容易错的是：

halo 跨 doc
compressed slot 跨 doc
query position 用全局 token 而不是 doc 内 position
cu_seqlens_cmp 直接用 padded_len // ratio

首版可以明确 loud fail：

packed meta is not None but packed support disabled

19. 最小可落地代码骨架

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __call__(self, query, key, value, attention_mask, x=None, qr=None):
        if self._cp_size() == 1:
            return self.origin_construct(query, key, value, attention_mask, x=x, qr=qr)

        ctx = self._build_cp_context(query, x)
        ratio = self._read_ratio()
        ori_kv = self._extract_ori_kv(key, value)

        ori_ext, halo_len = self._prepare_window_kv_halo(
            ori_kv, ctx, self._read_window_size()
        )

        cmp_main = None
        cmp_idx = None
        if ratio > 1:
            cmp_main = self._stage1_global_compressed_kv(
                self.parent.compressor, x, ctx
            )
        if ratio == 4:
            cmp_idx = self._stage1_global_compressed_kv(
                self.parent.indexer.compressor, x, ctx
            )

        if self.config.debug_force_smallop:
            return self.parent._construct_naive(
                query,
                ori_ext,
                x,
                qr,
                pos_offset=ctx.local_seq_start,
                cmp_kv_global=cmp_main,
                ori_halo_len=halo_len,
                cmp_idx_global=cmp_idx,
            )

        cmp_sparse_indices = None
        if ratio == 4:
            cmp_sparse_indices = self._run_global_indexer_topk(
                x, qr, cmp_idx, ctx
            )

        return self._run_sharedkv_attention(
            query=query,
            ori_kv_prepared=ori_ext,
            cmp_kv_global=cmp_main,
            cmp_sparse_indices=cmp_sparse_indices,
            ori_halo_len=halo_len,
            ctx=ctx,
        )

这个骨架的关键是：

Stage-1 在 HP 内部完成；
Stage-2 只用 local Q；
MF 不知道 halo/all-gather/select/top-k/KL loss 的细节。

20. 首版开发顺序建议

建议按这个顺序开工：

1. 新增 deepseek_v4_context_parallel.py 和 patch API。
2. 实现 CPContext 构建和输入校验。
3. 实现 window halo all-gather，写纯函数 UT。
4. 实现 ratio=128 global compressed KV。
5. 实现 ratio=4 compressor halo + select-and-pad。
6. 接 MF `_construct_naive` 做 small-op fallback，对齐 MF PR8311。
7. 加三个 PR717 v2/custom op wrapper 和 DistributedOp。
8. 切 fused sharedKV path。
9. 加 KL loss 和 backward 对齐。
10. 最后处理 packed sequence 和异步/P2P 优化。

这个顺序能把最大风险分开：

先证明 HP 的两阶段 CP 语义正确，
再证明 fused/custom op 接入正确，
最后做 packed 和性能。

yide12Collaborators

Jun 11 Close issue

yide12Collaborators

Jun 11 Comments：

HP 适配 DeepSeek V4 CP 的具体实现设计

参考代码：

MF PR8311：/home/wyd/code_review_deepseek_v4_cp/mindformers_pr8311_latest
HP 开发分支：/home/wyd/code_review_deepseek_v4_cp/hp_feature
MF 走读笔记：/home/wyd/code_review_deepseek_v4_cp/tmp.py

1. 总体结论

HP 需要新增的是一个 DeepSeek V4 CSA 两阶段 CP forward 能力，而不是一堆让 MF 自己拼的通信小函数。

MF 调用边界应该是：

hp.patch_deepseek_v4_compressed_sparse_attention_forward(
    csa_module,
    cp_mesh=cp_mesh,
    config=cp_config,
)

MF 负责选择哪些 CompressedSparseAttention module 需要 patch；HP 负责 patch 后完整接管 CP>1 的 attention forward。

核心 forward：

local Q / local ori_kv / local x / local qr
        |
        |-- window ori_kv halo
        |-- Stage-1: global compressed KV
        |-- ratio=4: global indexer top-k
        |
        v
local Q 在 [window halo + global compressed KV] 上算 attention
        |
        v
local output shard

最重要的工程判断：

已有 DSA ContextParallel style / DistributedOp 可以复用设计模式，
但不能直接等价替代 DeepSeek V4 两阶段 CP forward。

原因：

现有 DSAIndexerContextParallel / DSASparseAttentionContextParallel 是 op boundary hook，语义是“query shard，key/value gather”。
DeepSeek V4 CP 的关键不是普通 KV gather，而是：
- window KV halo；
- compressor halo；
- compressed KV all-gather；
- select-and-pad 去掉 PAD/DUP/无效 compressed slot；
- ratio=4 在 global compressed K 上做 top-k；
- sharedKV sparse attention 使用 prepared window KV + global compressed KV。

因此 HP 要做的是：

DeepSeekV4 CP forward 编排
    +
PR717 三个 custom op 的 DistributedOp/layout 语义

2. 当前 HP 分支能力盘点

当前 hp_feature 分支里已经有几类可复用基础。

2.1 可复用：ContextParallel 代码风格

已有文件：

hyper_parallel/core/context_parallel/context_parallel.py
hyper_parallel/core/context_parallel/dsa_context_parallel.py
hyper_parallel/core/context_parallel/async_dsa_context_parallel.py

可复用点：

ParallelStyle.apply(module, device_mesh) 的接口风格。
forward pre-hook / forward hook 的 patch 机制。
_ensure_1d(device_mesh)、DeviceMesh.get_group() 这类 CP mesh 获取方式。
_to_cp_dtensor(...)、Shard(seq_dim)、Replicate() 的 DTensor 标注方式。
async DSA 里 producer/consumer all-gather slot 的结构，后续可借鉴做异步 compressed KV gather。

不可直接复用点：

DSASparseAttentionContextParallel 会把 key/value 做 sequence replicate，这不等价于 DeepSeek V4 的 window halo。
DSAIndexerContextParallel 会把 key-side gather 成全局 K，但没有处理 compressor halo、select-and-pad、packed compressed cu_seqlens。
这些 style 不知道 MF CompressedSparseAttention 内部的 ratio/window/compressor/indexer 语义。

2.2 可复用：DistributedOp 写法

已有文件：

hyper_parallel/core/shard/ops/parallel_lightning_indexer.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_flash_attention.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_lightning_indexer_grad_kl_loss.py
hyper_parallel/core/shard/ops/yaml/lightning_indexer_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_flash_attention_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_lightning_indexer_grad_kl_loss_ops.yaml

可复用点：

preprocess -> infer_layout -> get_expand_impl 三段式。
对 BSND/TND layout 的参数归一化。
query 侧 Shard(seq)、key 侧 Replicate()、loss 输出 Partial(sum) 的布局推导。
TND packed seq lens 的调整模式。

缺口：

当前分支尚未看到 PR717 目标名字对应的 wrapper 和 DistributedOp：

npu_lightning_indexer_v2
npu_sparse_attn_sharedkv
npu_sparse_lightning_indexer_grad_kl_loss_v2

因此 DeepSeek V4 CP 需要新增这三个 op 的 HP wrapper / platform custom op / DistributedOp 规则。

2.3 可复用：通信接口

已有平台接口：

platform.differentiable_all_gather_concat(...)
platform.differentiable_reduce_scatter(...)
platform.all_gather_into_tensor(...)
platform.reduce_scatter_tensor(...)
platform.all_gather_single(...)
platform.reduce_scatter_single(...)

首版建议：

window halo:       all-gather boundary tail，然后本地 slice
compressor halo:   all-gather ratio 个边界 token，然后本地 slice
compressed KV:     all-gather local compressed slots
KL loss:           all-reduce / Partial(sum) reduce

P2P 可以作为后续优化，不进入首版正确性闭环。

3. 推荐新增文件

建议优先集中在 hyper_parallel/core/context_parallel/ 下，避免一开始拆太碎。

hyper_parallel/core/context_parallel/deepseek_v4_context_parallel.py

首版可以一个文件放完 public API、config、meta、forward class 和 small-op helper。等跑通后再拆：

hyper_parallel/core/context_parallel/deepseek_v4_cp_meta.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_indices.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_collectives.py
hyper_parallel/core/context_parallel/deepseek_v4_cp_forward.py

custom op 层新增：

hyper_parallel/custom_ops/experimental/experimental_ops.py
hyper_parallel/platform/mindspore/custom_ops/custom_ops.py
hyper_parallel/platform/mindspore/custom_ops/custom_op_impl.py
hyper_parallel/platform/mindspore/custom_ops/*.cc
hyper_parallel/core/shard/ops/parallel_npu_lightning_indexer_v2.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_attn_sharedkv.py
hyper_parallel/core/shard/ops/parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py
hyper_parallel/core/shard/ops/yaml/npu_lightning_indexer_v2_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_attn_sharedkv_ops.yaml
hyper_parallel/core/shard/ops/yaml/npu_sparse_lightning_indexer_grad_kl_loss_v2_ops.yaml

__init__ 导出：

hyper_parallel/core/context_parallel/__init__.py
hyper_parallel/__init__.py

新增 public symbol：

DeepSeekV4CPConfig
DeepSeekV4CPContext
DeepSeekV4CompressedSparseAttentionCPForward
patch_deepseek_v4_compressed_sparse_attention_forward

4. Public API 设计

4.1 Patch API

def patch_deepseek_v4_compressed_sparse_attention_forward(
        csa_module,
        *,
        cp_mesh,
        config,
):
    ...

含义：

csa_module：MF 选中的 CompressedSparseAttention Cell。
cp_mesh：HP DeviceMesh，首版只支持 1-D CP mesh。
config：HP DeepSeek V4 CP 配置。

原 csa_module 本体

不替换 Cell 对象，只 patch construct，保留 MF module 的参数注册和 state_dict。

4.2 Config

@dataclass
class DeepSeekV4CPConfig:
    enabled: bool = True
    mode: str = "colossal"
    input_layout: str = "SBNHD"       # MF CSA 边界: [S, B, N, D]
    op_layout: str = "BSND"           # NPU op 边界
    seq_dim: int = 0                  # MF CSA 边界 sequence dim
    op_seq_dim: int = 1               # BSND op sequence dim
    window_size: int | None = None    # None 表示从 csa_module 读取
    use_fused_ops: bool = True
    use_allgather_halo: bool = True
    strict: bool = True
    debug_force_smallop: bool = False

首版只支持：

mode == "colossal"
input_layout == "SBNHD"
op_layout == "BSND"

Ulysses/head shard 不进入 DeepSeek V4 CP 首版，因为 DSv4 custom ops 要求 head/head_dim/topk 维保持 replicated。

4.3 Runtime Context

DeepSeekV4CPContext 是 HP 在 forward 内部推导出来的运行时上下文：

@dataclass
class DeepSeekV4CPContext:
    cp_mesh: DeviceMesh
    cp_group: object
    cp_rank: int
    cp_size: int

    global_seq_len: int
    local_seq_start: int
    local_seq_end: int
    local_seq_len: int

    batch_size: int
    seq_dim: int
    op_layout: str

    compress_ratio: int
    window_size: int

    # packed sequence，可空
    cu_seqlens_q: object | None = None
    cu_seqlens_kv: object | None = None
    cu_seqlens_cmp: object | None = None
    cu_seqlens_padded: object | None = None
    cu_seqlens_real: object | None = None

DeepSeekV4CPContext 应由 HP 构建，不应要求 MF 直接填完整对象。

非 packed 等长序列只需要：

cp_mesh
query.shape[0]

HP 可推导：

cp_rank
cp_size
local_seq_len = query.shape[0]
global_seq_len = local_seq_len * cp_size
local_seq_start = cp_rank * local_seq_len
local_seq_end = local_seq_start + local_seq_len

packed/TND 场景需要额外知道文档边界：

actual_seq_len / cu_seqlens

这份信息不应通过 side-channel callback 传入 HP。它应该是正常 forward 数据流的一部分：

方案 A：MF 将 actual_seq_len/cu_seqlens 正常透传到 CompressedSparseAttention CP forward。
方案 B：如果 CSA 边界拿不到 actual_seq_len，则 MF 选择更高一层 DeepSeekV4HybridAttention 作为 HP CP hook 边界。

有了 actual_seq_len 后，HP 可继续推导：

每个 rank 的 global token range
每个 rank 覆盖哪些 doc
window halo 是否跨 doc
ratio=4 compressor halo 是否跨 doc
doc 内 query position
compressed cu_seqlens
select-and-pad gather_index

4.4 MF 调用方式

MF 侧推荐：

def apply_deepseek_v4_context_parallel(model, cp_mesh, config):
    for csa_module in model.iter_deepseek_v4_csa_modules():
        hp.patch_deepseek_v4_compressed_sparse_attention_forward(
            csa_module,
            cp_mesh=cp_mesh,
            config=config,
        )
    return model

HP 不提供：

hp.apply_deepseek_v4_context_parallel(model)

因为 HP 不应该遍历 MF 模型，也不应该根据 MF 结构猜 hook 位置。

5. Forward 替换对象

5.1 类结构

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __init__(self, parent, cp_mesh, config):
        self.parent = parent
        self.origin_construct = parent.construct
        self.cp_mesh = _ensure_1d(cp_mesh)
        self.config = config

    def __call__(self, query, key, value, attention_mask, x=None, qr=None, actual_seq_len=None):
        if not self.config.enabled or self._cp_size() == 1:
            return self.origin_construct(
                query, key, value, attention_mask, x=x, qr=qr
            )

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

Patch 函数：

def patch_deepseek_v4_compressed_sparse_attention_forward(...):
    cp_forward = DeepSeekV4CompressedSparseAttentionCPForward(...)
    csa_module._hp_origin_construct = csa_module.construct
    csa_module._hp_deepseek_v4_cp_forward = cp_forward

    def construct(query, key, value, attention_mask, x=None, qr=None):
        return cp_forward(query, key, value, attention_mask, x=x, qr=qr)

    csa_module.construct = construct
    return csa_module

5.2 入口输入 shape

MF CompressedSparseAttention.construct 边界是 seq-first：

query: [s_local, B, N, D]
key:   [s_local, B, 1, D] 或 [s_local, B, D]
value: [s_local, B, 1, D] 或 [s_local, B, D]
x:     [s_local, B, H]
qr:    [s_local, B, q_lora_rank]

HP 入口必须强校验：

query.shape[0] == cp_ctx.local_seq_len
key.shape[0]   == cp_ctx.local_seq_len
x.shape[0]     == cp_ctx.local_seq_len
qr.shape[0]    == cp_ctx.local_seq_len

ratio 约束：

ratio=0:   x/qr 可不参与 compressed 分支
ratio=128: x 必须存在
ratio=4:   x 和 qr 必须存在，且 parent.indexer 必须存在

6. CP Forward 主流程

主流程伪代码：

def _construct_cp(self, query, key, value, attention_mask, x, qr, cp_ctx):
    parent = self.parent
    ratio = self._read_compress_ratio(parent)

    ori_kv = self._extract_shared_kv(key, value)

    ori_ext, ori_halo_len = self._prepare_window_kv_halo(
        ori_kv, cp_ctx, window_size=parent.window_size
    )

    cmp_main = None
    cmp_idx = None
    if ratio and ratio > 1:
        cmp_main = self._stage1_global_compressed_kv(
            parent.compressor, x, cp_ctx, purpose="main"
        )

    if ratio == 4 and parent.indexer is not None:
        cmp_idx = self._stage1_global_compressed_kv(
            parent.indexer.compressor, x, cp_ctx, purpose="indexer"
        )

    if self.config.debug_force_smallop:
        return self._stage2_smallop(
            query, ori_ext, x, qr,
            cmp_main, cmp_idx,
            ori_halo_len, cp_ctx,
        )

    return self._stage2_fused_sharedkv(
        query, ori_ext, x, qr,
        cmp_main, cmp_idx,
        ori_halo_len, cp_ctx,
    )

建议首个可验证版本先实现：

debug_force_smallop=True

7. Window KV Halo

7.1 功能

window halo 用来解决 rank 边界处 query 看不到左侧 KV 的问题。

例子：

CP = 2
s_local = 4
window_size = 4

rank0: 0 1 2 3
rank1: 4 5 6 7

rank1 的 token 4 要看：

1 2 3 4

所以 rank1 需要从 rank0 拿：

1 2 3

这段就是 window KV halo。

7.2 首版通信

首版用 all-gather boundary tail：

hlen = min(window_size - 1, s_local)
local_tail = ori_kv[-hlen:]                 # [hlen, B, D]
all_tails = all_gather_seq(local_tail)      # [cp_size * hlen, B, D]
left_halo = slice_left_context(all_tails, cp_rank, hlen, window_size)
ori_ext = cat([left_halo, ori_kv], dim=0)

shape：

ori_kv:    [s_local, B, D]
local_tail:[hlen, B, D]
all_tails: [cp_size * hlen, B, D]
left_halo: [halo_len, B, D]
ori_ext:   [halo_len + s_local, B, D]

hlen 是每个 rank 贡献多少 boundary，不是当前 rank 最终拿多少 halo：

hlen = min(window_size - 1, s_local)

当 window_size - 1 > s_local 时，当前 rank 的 left_halo 可能来自多个前序 rank；all-gather 后本地 slice 可以统一处理。

7.3 函数建议

def _prepare_window_kv_halo(ori_kv, cp_ctx, window_size):
    hlen = min(window_size - 1, cp_ctx.local_seq_len)
    if hlen <= 0 or cp_ctx.cp_rank == 0:
        ...
    all_tails = _all_gather_seq(ori_kv[-hlen:], cp_ctx)
    left_halo = _assemble_left_halo_from_rank_tails(
        all_tails,
        cp_rank=cp_ctx.cp_rank,
        hlen=hlen,
        need=window_size - 1,
    )
    return cat([left_halo, ori_kv], dim=0), left_halo.shape[0]

注意 packed sequence：

halo 不能跨文档。

首版非 packed 可先按连续全局 token slice；packed 版本必须按 cu_seqlens_real 对每个 doc 单独组 halo。

8. Stage-1：Global Compressed KV

Stage-1 的唯一目标：

让每个 rank 都拿到 CP=1 语义下的 global compressed KV。

Q 不在 Stage-1 压缩。x 分两路：

x -> attention projection -> Q / ori_kv
x -> compressor           -> compressed KV

8.1 ratio=0

无 compressed KV：

cmp_main = None
cmp_idx = None

只走 window halo + sharedKV attention。

8.2 ratio=128

非重叠压缩。

流程：

local_slots = compressor(
    x,
    rope_pos_offset=cp_ctx.local_compressed_start * ratio,
)
gathered = all_gather_seq(local_slots)
cmp_global = select_and_pad(gathered, gather_index)

非 packed 且 s_local % 128 == 0 时：

local_slots: [s_local / 128, B, D]
gathered:    [S / 128, B, D]
cmp_global:  [S / 128, B, D]

需要注意因果可见性：

query position p 只能看 compressed slot c < (p + 1) // ratio

例如 p=300，ratio=128：

visible slots: 0, 1
slot 2 压缩 256..383，包含未来 301..383，不能看。

8.3 ratio=4

重叠压缩，需要 compressor halo。

MF 当前语义可以抽象为：

1. 当前 rank 拿左侧 ratio 个 hidden states 作为 compressor halo
2. compressor 在 [halo_x, local_x] 上压缩
3. 每个 rank 生成等长 local slots
4. rank0 的第一个 slot 是 PAD
5. rank>0 的第一个 slot 是跨边界 DUP
6. all-gather 后用 select-and-pad 去掉 PAD/DUP/无效 slot

首版实现建议继续使用 all-gather boundary：

compress_hlen = ratio
x_tail = x[-compress_hlen:]
all_x_tails = all_gather_seq(x_tail)
halo_x = slice_left_context(all_x_tails, cp_rank, compress_hlen, need=ratio)
local_slots = stage1_compress_with_halo(compressor, x, halo_x, cp_ctx, overlap=True)
gathered = all_gather_seq(local_slots)
cmp_global = select_and_pad(gathered, gather_index)

main compressed KV 和 indexer compressed K 都要做：

cmp_main = stage1_global(parent.compressor)
cmp_idx  = stage1_global(parent.indexer.compressor)

shape：

x:             [s_local, B, H]
halo_x:        [ratio, B, H] 或更短
local_slots:   [local_slot_count_with_pad_dup, B, D]
gathered:      [cp_size * local_slot_count_with_pad_dup, B, D]
cmp_global:    [C, B, D]

cmp_idx:       [C, B, index_head_dim]

8.4 select-and-pad

ratio=4 的 overlap 会产生 PAD/DUP；packed sequence 还会产生文档 padding 无效块。因此 all-gather 后不能直接使用 gathered。

需要一个 index：

gather_index: [C]

做：

padded = cat([gathered, zero_sentinel], dim=0)
cmp_global = index_select(padded, dim=0, index=gather_index)

其中：

valid slot -> 指向 gathered 的真实位置
invalid    -> 指向 zero_sentinel

函数建议：

def build_compressed_gather_index(cp_ctx, ratio, overlap):
    ...

def select_and_pad_compressed(gathered, gather_index):
    ...

缓存建议：

key = (cp_size, global_seq_len, local_seq_len, ratio, packed_signature)

首版非 packed 可只支持：

S % (cp_size * ratio) == 0

packed 支持后再加入 cu_seqlens_real/padded。

9. Stage-2：Local Q Attention

Stage-2 输入：

query:     [s_local, B, N, D]
ori_ext:   [halo_len + s_local, B, D]
cmp_main:  [C, B, D] 或 None
cmp_idx:   [C, B, index_head_dim]，ratio=4 only

Stage-2 输出：

output: [s_local, B, N * D]

9.1 Small-op fallback 路径

用于首版对齐 MF PR8311。

如果 MF module 有 _construct_naive：

out = parent._construct_naive(
    query,
    ori_ext,
    x,
    qr,
    pos_offset=cp_ctx.local_seq_start,
    cmp_kv_global=cmp_main,
    ori_halo_len=halo_len,
    cmp_idx_global=cmp_idx,
)

语义：

kv_full = [ori_ext, cmp_main]
window_idxs 指向 kv_full 前半段
compressed_idxs 指向 kv_full 后半段
topk_idxs = [window_idxs, compressed_idxs]
unfused_compressed_sparse_attn(query, kv_full, topk_idxs)

这个 fallback 很适合做第一阶段 HP 验证：

HP 通信/metadata 正确性
MF CP=1 vs HP CP>1 数值对齐

但最终 fused 目标不应长期依赖 MF 私有 _construct_naive。

9.2 Fused sharedKV 路径

最终目标是调用 PR717 的：

npu_sparse_attn_sharedkv

输入准备：

q_local:        local Q
ori_kv_prepared: [left_halo, local_ori_kv]
cmp_kv_global: global compressed KV
cmp_sparse_indices:
  ratio=0:   None
  ratio=128: None，op 内部按 cmp_ratio/causal 处理 dense compressed visibility
  ratio=4:   global top-k indices from npu_lightning_indexer_v2

调用示意：

out = npu_sparse_attn_sharedkv(
    q_local_dt,
    ori_kv_prepared_dt,
    cmp_kv_global_dt,
    cmp_sparse_indices_dt,
    sinks=parent.attn_sink,
    scale_value=parent.softmax_scale,
    actual_seq_qlen=cp_ctx.cu_seqlens_q,
    actual_seq_klen=cp_ctx.cu_seqlens_kv,
    actual_seq_cmp_klen=cp_ctx.cu_seqlens_cmp,
    layout=cp_ctx.op_layout,
    cmp_ratio=ratio or None,
    ori_win_left=parent.ori_win_left,
    ori_win_right=parent.ori_win_right,
    local_seq_start=cp_ctx.local_seq_start,
    ori_halo_len=halo_len,
)

注意：

ori_kv_prepared 不是普通全局 replicated KV。
它是 HP 已经拼好的 [left_halo, local_ori_kv]。
DistributedOp 不能再把它当普通 key/value 做 all-gather。

因此 parallel_npu_sparse_attn_sharedkv.py 的布局规则要有 prepared window KV 语义。

10. ratio=4 Global Top-k

ratio=4 不能只在本 rank compressed K 上 top-k，必须使用 global indexer compressed K：

local q_idx / local weights
    +
global cmp_idx
    +
global causal mask / actual seq meta
    |
    v
global top-k compressed indices

10.1 q_idx / weights

复用 MF indexer 的前半段逻辑：

q_idx, _, w_idx = parent.indexer.forward_before_topk(
    stop_gradient(x),
    stop_gradient(qr),
    query_pos_offset=cp_ctx.local_seq_start,
)

shape：

q_idx: [s_local, B, index_n_heads, index_head_dim]
w_idx: [s_local, B, index_n_heads]

如果 fused op 需要 BSND：

q_idx -> [B, s_local, index_n_heads, index_head_dim]
w_idx -> [B, s_local, index_n_heads]
cmp_idx_global -> [B, C, index_n_heads or 1, index_head_dim]  # 以算子真实签名为准

10.2 调用 `npu_lightning_indexer_v2`

topk_indices, topk_scores = npu_lightning_indexer_v2(
    q_idx_dt,
    cmp_idx_global_dt,
    w_idx_dt,
    actual_seq_qlen=cp_ctx.cu_seqlens_q,
    actual_seq_klen=cp_ctx.cu_seqlens_cmp,
    layout=cp_ctx.op_layout,
    sparse_count=parent.indexer.index_topk,
    sparse_mode=3,
    cmp_ratio=4,
    query_pos_offset=cp_ctx.local_seq_start,
)

布局：

q_idx:          Shard(seq)
w_idx:          Shard(seq)
cmp_idx_global: Replicate()
topk_indices:   Shard(seq)
topk_scores:    Shard(seq)

输出 indices 是 compressed slot id，不是 kv_full 后半段 id。small-op fallback 里需要加 offset；fused sharedKV op 通常直接接收 compressed slot id。

10.3 Causal mask

因果规则：

compressed slot c 对 query global position p 可见 iff c < (p + 1) // ratio

packed sequence 下要按每个 doc 内位置算，不能用跨 doc 全局 token 直接除。

11. KL Loss

ratio=4 训练态需要 indexer KL loss。

目标 op：

npu_sparse_lightning_indexer_grad_kl_loss_v2

输入布局：

query:         Shard(seq)
key:           Replicate()  # main compressed/global key 或 sharedKV key，以 op 签名为准
query_index:   Shard(seq)
key_index:     Replicate()  # cmp_idx_global
weights:       Shard(seq)
sparse_indices:Shard(seq)
softmax_max:   Shard(seq)
softmax_sum:   Shard(seq)

输出：

d_query_index: Shard(seq)
d_key_index:   Replicate() 或 Partial(sum) 后 reduce
d_weights:     Shard(seq)
loss:          Partial(sum) -> all-reduce 后全局标量

HP 需要保证：

CP>1 的 KL loss 数值等价 CP=1。

工程上可沿用现有 DSAIndexerLossContextParallel 对 d_key_index 和 loss partial 的处理思路，但 op 名称、参数和 v2 语义要重新实现。

12. DistributedOp 规则

12.1 `parallel_npu_lightning_indexer_v2.py`

约束：

query_index: [B, S1, N1, D]，S1 可 Shard，N1/D replicated
key_index:   [B, S2, N2, D]，S2/N2/D replicated
weights:     [B, S1, N1]，B/S1 与 query_index 一致，N1 replicated

输出：

topk_indices: [B, S1, N2, sparse_count]
topk_scores:  [B, S1, N2, sparse_count]

布局：

B/S1 继承 query_index
N2/sparse_count replicated

CP expand：

key_index 已由 Stage-1 变成 global replicate，不再按普通 causal window 截 key。
只需要根据 local_seq_start / actual_seq_qlen / cmp_ratio 处理 causal visibility。

12.2 `parallel_npu_sparse_attn_sharedkv.py`

输入角色：

query:              Shard(seq)
ori_kv_prepared:    local prepared window KV with halo
cmp_kv_global:      Replicate()
cmp_sparse_indices: Shard(seq)，ratio=4 only

输出：

attention_out: Shard(seq)
softmax/lse:   Shard(seq) 或 op 定义布局

重点：

ori_kv_prepared 不能由 DistributedOp 自动 all-gather。
halo 已由 DeepSeekV4 CP forward 准备好。

建议在 wrapper 参数中显式标识：

ori_kv_is_prepared=True
ori_halo_len=halo_len
local_seq_start=cp_ctx.local_seq_start

如果用户直接拿普通 DTensor KV 调这个 op，应报错或要求走普通 DSA/SFA op。

12.3 `parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py`

复用现有 sparse KL loss DistributedOp 的大方向：

d_query_index: 继承 query_index layout
d_key_index:   key_index layout + CP Partial(sum)
d_weights:     继承 weights layout
loss:          scalar Partial(sum) over CP/DP

需要额外处理：

key_index 是 global compressed K，不是普通 full sequence K。
actual_seq_klen 应是 compressed cu_seqlens。

13. CP Meta 和 Packed Sequence

非 packed 首版：

global_seq_len = local_seq_len * cp_size
local_seq_start = cp_rank * local_seq_len
local_seq_end = local_seq_start + local_seq_len
cu_seqlens_q = None
cu_seqlens_cmp = None

packed 版本必须处理：

1. window halo 不跨文档
2. ratio=4 compressor halo 不跨文档
3. compressed cu_seqlens 按 doc 内 real length floor(real_len / ratio)
4. causal visible slot 按 doc 内 query position 计算
5. select-and-pad index 按 doc 生成

packed meta 的来源规则：

actual_seq_len / cu_seqlens 必须走正常 forward 数据流。

也就是说：

1. 如果 CompressedSparseAttention 边界能拿到 actual_seq_len，
   HP patch 这个边界。

2. 如果 CompressedSparseAttention 边界拿不到 actual_seq_len，
   MF 应选择更高一层 DeepSeekV4HybridAttention 作为 HP CP hook 边界，
   由这个边界把 actual_seq_len 和 query/key/x/qr 一起交给 HP。

HP CP forward 拿到 actual_seq_len/cu_seqlens 后，再自己推导：

local/global token range
doc 内 position
window halo 边界
compressor halo 边界
compressed cu_seqlens
select-and-pad gather_index

这里不引入 meta_provider、全局变量或 side-channel callback。完整
DeepSeekV4CPContext 是 HP 内部派生物，不由 MF 填。

14. 通信实现细节

14.1 `_all_gather_seq`

内部统一封装：

def _all_gather_seq(tensor, cp_ctx):
    return platform.differentiable_all_gather_concat(
        tensor,
        group=cp_ctx.cp_group,
        concat_size=cp_ctx.cp_size,
        concat_dim=0,
    )

MindSpore 平台也可以直接使用：

comm_func.all_gather_into_tensor(None, tensor.contiguous(), group=group)

但 HP 内部优先走 platform 抽象，方便 Torch/MindSpore 双后端和 autograd。

14.2 为什么首版不用 P2P

all-gather 逻辑：

每个 rank 发一样大小的 boundary
每个 rank 收所有 boundary
每个 rank 本地 slice

P2P 逻辑：

每个 rank 要算自己从哪些前序 rank 收、给哪些后续 rank 发、每个 peer 发多少、backward 怎么还梯度。

DeepSeek V4 CP 首版优先正确性，all-gather 更适合。

14.3 backward

首版依赖 differentiable all-gather 的 VJP：

forward all-gather
backward reduce-scatter / slice grad accumulate

select-and-pad 的 backward：

index_select forward
scatter_add backward

如果 MindSpore 小算子没有自动满足，需要实现自定义 grad 或使用已有 differentiable gather/scatter 组合。

15. 与 MF PR8311 的逐项映射

MF PR8311 能力	HP 目标实现位置
`two_stage_cp_naive_attention` 总控	`DeepSeekV4CompressedSparseAttentionCPForward._construct_cp`
window KV halo all-gather	`_prepare_window_kv_halo`
`_gather_compressed(cell.compressor)`	`_stage1_global_compressed_kv(parent.compressor)`
`_gather_compressed(cell.indexer.compressor)`	`_stage1_global_compressed_kv(parent.indexer.compressor)`
`stage1_compress_with_halo`	`_compress_local_slots_with_halo`
`build_select_pad_gather_index`	`build_compressed_gather_index`
`select_and_pad_smallop`	`select_and_pad_compressed`
`_construct_naive(... cmp_kv_global ...)`	small-op fallback / fused sharedKV path
ratio=4 `fused_qk_topk_naive`	`npu_lightning_indexer_v2`
`compute_dsa_indexer_loss`	`npu_sparse_lightning_indexer_grad_kl_loss_v2`

16. 开发里程碑

M1：HP small-op CP forward 跑通

范围：

新增 patch API。
新增 DeepSeekV4CPContext。
实现 window halo all-gather。
实现 ratio=128 global compressed KV。
实现 ratio=4 compressor halo + all-gather + select-and-pad。
Stage-2 先调用 MF _construct_naive 或 HP 等价 small-op。

验收：

ratio=0/4/128
CP=1 vs CP=2/4
forward output 对齐

M2：接入 PR717 v2/custom sharedKV ops

范围：

新增三个 experimental op wrapper。
MindSpore platform custom op 接入。
三个 DistributedOp + YAML 注册。
Fused Stage-2 使用 npu_sparse_attn_sharedkv。
ratio=4 top-k 使用 npu_lightning_indexer_v2。

验收：

small-op fallback vs fused path 对齐
CP=1 fused vs CP>1 fused 对齐

M3：KL loss 和反向

范围：

接入 npu_sparse_lightning_indexer_grad_kl_loss_v2。
loss all-reduce / Partial(sum)。
d_key_index CP partial reduce。
window halo / compressed gather / select-pad backward 验证。

验收：

loss 对齐
主要参数梯度对齐

M4：packed sequence

范围：

meta provider 支持 packed cu_seqlens。
window halo 不跨 doc。
compressor halo 不跨 doc。
compressed cu_seqlens per-doc。
ratio=4 select-and-pad per-doc。

验收：

多 doc packed batch
不同长度
rank 边界切在 doc 内 / doc 边界上 / padding 区

M5：异步和性能优化

范围：

window halo P2P 可选实现。
compressed KV all-gather async overlap。
indexer compressed K gather 与 q_idx 计算 overlap。
gather_index cache。

验收：

数值不变
通信时间下降
无额外峰值内存爆炸

17. 测试计划

17.1 纯函数 UT

新增：

tests/ut/core/context_parallel/test_deepseek_v4_cp_indices.py

覆盖：

hlen = min(window_size - 1, s_local)。
window halo 多 rank slice。
window_size - 1 < s_local 只需要前一个 rank。
window_size - 1 > s_local 需要多个前序 rank。
ratio=128 causal visible compressed slots。
ratio=4 select-and-pad index 去 PAD/DUP。

17.2 CP forward small-op UT/ST

新增：

tests/mindspore/st/context_parallel/deepseek_v4_cp_smallop.py
tests/mindspore/st/context_parallel/test_deepseek_v4_cp_smallop_msrun.py

覆盖：

ratio=0, window only
ratio=128, dense compressed
ratio=4, indexer top-k
CP=1/2/4
window 跨 0/1/多 rank

对齐标准：

CP=1 全量输入 reference
CP>1 本地输出 all-gather 后与 reference 对齐

17.3 DistributedOp UT

新增：

tests/ut/core/shard/test_parallel_npu_lightning_indexer_v2.py
tests/ut/core/shard/test_parallel_npu_sparse_attn_sharedkv.py
tests/ut/core/shard/test_parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py

覆盖：

layout inference。
非法 head/head_dim/topk sharding 报错。
query Shard(seq) 输出仍 Shard(seq)。
key/compressed K 必须 replicated。
KL loss 输出 partial reduce。

17.4 Fused ST

新增：

tests/mindspore/st/context_parallel/deepseek_v4_cp_fused.py
tests/mindspore/st/context_parallel/test_deepseek_v4_cp_fused_msrun.py

覆盖：

small-op fallback vs fused
CP=1 fused vs CP>1 fused
forward + backward

18. 风险和约束

18.1 不要把通信原语暴露给 MF

不推荐：

halo = hp.exchange_window_halo(...)
cmp = hp.all_gather_compressed_kv(...)
idx = hp.build_select_pad(...)

18.2 不要只加 DistributedOp 就结束

local compressed KV 已经错了，
底层 op 再聪明也补不回 CP=1 语义。

18.3 prepared window KV 不是普通 replicated KV

DeepSeek V4 sharedKV op 的 ori_kv 输入在 CP forward 里已经变成：

[left_halo, local_ori_kv]

它不是：

all-gather 后的 full ori_kv

DistributedOp 规则必须区分这个语义。

18.4 packed sequence 是单独难点

非 packed 跑通不代表 packed 正确。

packed 下最容易错的是：

halo 跨 doc
compressed slot 跨 doc
query position 用全局 token 而不是 doc 内 position
cu_seqlens_cmp 直接用 padded_len // ratio

首版可以明确 loud fail：

packed meta is not None but packed support disabled

19. 最小可落地代码骨架

class DeepSeekV4CompressedSparseAttentionCPForward:
    def __call__(self, query, key, value, attention_mask, x=None, qr=None):
        if self._cp_size() == 1:
            return self.origin_construct(query, key, value, attention_mask, x=x, qr=qr)

        ctx = self._build_cp_context(query, x)
        ratio = self._read_ratio()
        ori_kv = self._extract_ori_kv(key, value)

        ori_ext, halo_len = self._prepare_window_kv_halo(
            ori_kv, ctx, self._read_window_size()
        )

        cmp_main = None
        cmp_idx = None
        if ratio > 1:
            cmp_main = self._stage1_global_compressed_kv(
                self.parent.compressor, x, ctx
            )
        if ratio == 4:
            cmp_idx = self._stage1_global_compressed_kv(
                self.parent.indexer.compressor, x, ctx
            )

        if self.config.debug_force_smallop:
            return self.parent._construct_naive(
                query,
                ori_ext,
                x,
                qr,
                pos_offset=ctx.local_seq_start,
                cmp_kv_global=cmp_main,
                ori_halo_len=halo_len,
                cmp_idx_global=cmp_idx,
            )

        cmp_sparse_indices = None
        if ratio == 4:
            cmp_sparse_indices = self._run_global_indexer_topk(
                x, qr, cmp_idx, ctx
            )

        return self._run_sharedkv_attention(
            query=query,
            ori_kv_prepared=ori_ext,
            cmp_kv_global=cmp_main,
            cmp_sparse_indices=cmp_sparse_indices,
            ori_halo_len=halo_len,
            ctx=ctx,
        )

这个骨架的关键是：

Stage-1 在 HP 内部完成；
Stage-2 只用 local Q；
MF 不知道 halo/all-gather/select/top-k/KL loss 的细节。

20. 首版开发顺序建议

建议按这个顺序开工：

1. 新增 deepseek_v4_context_parallel.py 和 patch API。
2. 实现 CPContext 构建和输入校验。
3. 实现 window halo all-gather，写纯函数 UT。
4. 实现 ratio=128 global compressed KV。
5. 实现 ratio=4 compressor halo + select-and-pad。
6. 接 MF `_construct_naive` 做 small-op fallback，对齐 MF PR8311。
7. 加三个 PR717 v2/custom op wrapper 和 DistributedOp。
8. 切 fused sharedKV path。
9. 加 KL loss 和 backward 对齐。
10. 最后处理 packed sequence 和异步/P2P 优化。

这个顺序能把最大风险分开：

先证明 HP 的两阶段 CP 语义正确，
再证明 fused/custom op 接入正确，
最后做 packed 和性能。

yide12Collaborators

Jun 11 Reopen issue

Hyper-Parallel 支持 MindFormers DeepSeek V4 CP 并行方案

背景与结论

职责边界

MF 负责

HP 负责

MF 与 HP 的接口契约

风格对齐

前向替换机制

替换前向不是包一层原 forward

HP 是否要完整重写 attention

CP forward 内部仍要触发 DistributedOp

输入 Sequence Sharding

HP 侧新增模块建议

二阶段 CP 总体算法

Stage-1：构建 compressed KV

ratio=0

ratio=128

ratio=4

Stage-2：本地 Q shared-KV sparse attention

ratio=0

ratio=128

ratio=4

三个自定义算子的分布式语义

npu_lightning_indexer_v2

npu_sparse_attn_sharedkv

npu_sparse_lightning_indexer_grad_kl_loss_v2

Packed Sequence / TND 元数据

实施阶段

Phase 1：算子接入和分布式注册

Phase 2：MF/HP 能力接口

Phase 3：ratio0 / ratio128

Phase 4：ratio4 CSA 两段式前半段

Phase 5：ratio4 CSA 两段式后半段

Phase 6：TND / packed sequence

Phase 7：TP+CP 和性能优化

最小测试矩阵

风险点

一句话方案

MindFormers 两段式 DeepSeek V4 CP 实现拆解

1. 先给结论

2. MF 如何决定替换哪个 forward

3. MF 如何传 CP meta

4. MF 如何保证输入已经是 local sequence

5. CP forward 的实际入口

6. Stage-1：compressed KV 全局化

6.1 ratio4 CSA：compressor halo

6.2 ratio128 HCA：非重叠 compressed KV all-gather

6.3 当前 packed meta 缺口

7. Stage-2：local Q attention

7.1 window KV halo

7.2 调用 _construct_naive

7.3 HCA global visibility

7.4 CSA global top-k

8. KL loss 当前是怎么做的

9. 反向通信对称性

10. MF 测试如何证明 CP forward 正确

11. 和你列出的 HP 第二类能力逐项对照

12. 对 HP 接口的直接启示

MindFormers 两段式 DeepSeek V4 CP 实现拆解

1. 先给结论

2. MF 如何决定替换哪个 forward

3. MF 如何传 CP meta

4. MF 如何保证输入已经是 local sequence

5. CP forward 的实际入口

6. Stage-1：compressed KV 全局化

6.1 ratio4 CSA：compressor halo

6.2 ratio128 HCA：非重叠 compressed KV all-gather

6.3 当前 packed meta 缺口

7. Stage-2：local Q attention

7.1 window KV halo

7.2 调用 _construct_naive

7.3 HCA global visibility

7.4 CSA global top-k

8. KL loss 当前是怎么做的

9. 反向通信对称性

10. MF 测试如何证明 CP forward 正确

11. 和你列出的 HP 第二类能力逐项对照

12. 对 HP 接口的直接启示

HP 适配 DeepSeek V4 CP 的具体实现设计

1. 总体结论

7.2 调用 `_construct_naive`

7.2 调用 `_construct_naive`

10.2 调用 `npu_lightning_indexer_v2`

12.1 `parallel_npu_lightning_indexer_v2.py`

12.2 `parallel_npu_sparse_attn_sharedkv.py`

12.3 `parallel_npu_sparse_lightning_indexer_grad_kl_loss_v2.py`

14.1 `_all_gather_seq`