[RFC] DSA2 Sequence Replicate Cache 跨边界归属与 MF/HP 职责设计

不建议直接把缓存逻辑塞进 DSASparseAttentionContextParallel。虽然当前只用于 DSA2，但被共享的数据跨越了两个独立模块，Sparse Attention 只是第一个消费者，不是完整生命周期的所有者。

当前数据生命周期

Sparse Attention pre-hook
    ├─ cache.begin()
    ├─ AllGather main_kv/key_rope
    └─ 将结果存入 cache
             │
             ▼
Sparse Attention forward
             │
             ▼
Indexer Loss pre-hook
    ├─ 从 cache 复用 main_kv/key_rope
    └─ cache.clear()

对应代码：

• Sparse Attention 创建缓存内容：hyper_parallel/core/context_parallel/dsa_context_parallel.py:394
• Indexer Loss 消费并清理缓存：hyper_parallel/core/context_parallel/dsa_context_parallel.py:817

因此缓存的实际所有者应该覆盖：

Sparse Attention + Indexer Loss

而不是只属于 DSASparseAttentionContextParallel。

直接放进 Sparse Attention 的问题

如果缓存作为 DSASparseAttentionContextParallel 的成员：

sparse_style.cache = {...}

Indexer Loss 只能通过以下方式复用：

indexer_loss_style.sparse_attention_style = sparse_style

这会导致：

• Indexer Loss 依赖另一个 ParallelStyle 的内部状态；
• 两个独立 hook 的执行顺序变成隐式约束；
• Sparse Attention 需要管理 Indexer Loss 结束后的清理；
• 后续一个 Sparse Attention 对应多个 loss consumer 时难以扩展；
• 单独测试 Indexer Loss style 时必须构造假的 Sparse Attention style；
• 容易因异常路径未清理而引用上一轮 forward 的 Tensor。

更重要的是，该类不仅是普通字典，还负责：

• 校验相同 slot 的 shape、dtype、CP group 和 sequence dim；
• detach 唯一的 AllGather 底层存储；
• 给每个消费者建立独立 Gradient Bridge；
• 保持各分支独立 ReduceScatter；
• 控制 begin → reuse → clear 生命周期。

这些职责属于“跨边界通信共享上下文”，不属于 Sparse Attention 算子本身。

MF 与 HP 的职责边界

use_sparse_loss、async_enabled 以及 DSA1/DSA2 阶段判断属于 MindFormers 配置语义，Hyper-Parallel 当前并没有、也不应该增加 use_sparse_loss 这类模型阶段参数。

因此正确的职责拆分是：

MindFormers
    ├─ 判断当前是否为同步 DSA2
    ├─ 决定是否创建共享上下文
    └─ 将同一个共享上下文注入两个 HP style

Hyper-Parallel
    ├─ 实现 Sequence AllGather 结果复用
    ├─ 管理 begin/reuse/clear 生命周期
    ├─ 校验 shape/dtype/CP group/seq_dim
    └─ 为每个消费者建立独立 Gradient Bridge

当前 MindFormers 的启用条件应继续保留：

shared_replicate_cache = (
    HPDSASequenceReplicateCache()
    if use_sparse_loss and not async_enabled
    else None
)

attention_style = ...(
    shared_replicate_cache=shared_replicate_cache
)
indexer_loss_style = ...(
    shared_replicate_cache=shared_replicate_cache
)

HP 不需要知道 use_sparse_loss，只接收 MF 已经完成阶段判断后传入的共享上下文。

是否需要进一步封装

如果希望减少 MF 对 HP 内部实现的了解，可以在 HP 提供一个不包含 DSA 阶段语义的通用组合接口，例如：

DSASparseBoundaryGroup
    ├─ DSASparseAttentionContextParallel
    ├─ DSAIndexerLossContextParallel
    └─ DSASequenceReplicateCache

该接口只接收通用的 enable_shared_sequence_replicate，不接收 use_sparse_loss：

class DSASparseBoundaryGroup:
    def __init__(self, ..., enable_shared_sequence_replicate=False):
        cache = (
            DSASequenceReplicateCache()
            if enable_shared_sequence_replicate
            else None
        )

        self.attention_style = DSASparseAttentionContextParallel(
            shared_replicate_cache=cache,
        )
        self.indexer_loss_style = DSAIndexerLossContextParallel(
            shared_replicate_cache=cache,
)

MF 仍负责映射：

enable_shared_sequence_replicate = use_sparse_loss and not async_enabled

但当前只有两个消费者，现有显式注入已经足够清晰。为隐藏一次对象创建而新增组合 API，收益可能小于接口复杂度。因此本 PR 不一定需要继续封装。

结论

推荐程度：

1. 推荐：保留独立的 DSASequenceReplicateCache；MF 根据 use_sparse_loss and not async_enabled 决定是否创建，HP 负责通信和梯度语义。
2. 可选优化：如果未来消费者继续增加，可在 HP 增加不感知 DSA1/DSA2 的通用组合对象，由 MF 传入 enable_shared_sequence_replicate。
3. 不推荐：把缓存直接作为 DSASparseAttentionContextParallel 的内部成员，再让 Indexer Loss 依赖 Attention 对象。

“当前只服务 DSA2”是 MF 的启用范围；“缓存跨越两个消费者”是 HP 的通信机制。两者应保持分层，HP 不应新增 use_sparse_loss 参数。

关联 PR：https://gitcode.com/mindspore/hyper-parallel/pull/1282

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 356
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/356

Start with hyper_parallel/core/context_parallel/dsa_context_parallel.py at lines 394 and 817, then review the associated PR 1282. Trace how Sparse Attention creates the cache and how Indexer Loss consumes and clears it. Done means the shared context has an explicit lifecycle and independent ownership without adding MindFormers-specific stage parameters to Hyper-Parallel.