任务4：新流程 CP 多通信模式迁移设计方案

1. 结论

当前 refactor_hp/hyper_parallel/components/distributed 中的 CP 实现，实际上只有一种通信算法：

本地 Q + K/V all-gather + 本地 attention

旧流程 hp/hyper_parallel/core/context_parallel/ 已经具备：

• Pure Colossal：K/V all-gather；
• Pure Ulysses：Q/K/V sequence-to-head all-to-all；
• Hybrid：Ulysses 子组 all-to-all + Colossal 子组 K/V all-gather；
• 上述模式的同步和部分异步重叠；
• Colossal Head-Tail 负载均衡；
• DSA Indexer / SparseAttention / Loss 的同步、异步 CP。

迁移到新流程时，不能简单增加 sdpa_qkv_ulysses、sdpa_hf_hybrid_async 等 wrapper。这样会把“网络调用方式 × attention primitive × CP 算法 × 同异步”做成笛卡尔积，注册表会快速膨胀。

推荐把 CP 拆成两个正交层次：

Attention Adapter
  负责：在哪里拿到 Q/K/V、怎样调用原始 attention、怎样恢复输出
  类型：sdpa_qkv / sdpa_hf / flex_qkv / flex_hf / DSA / custom

CP Strategy
  负责：Q/K/V 做什么通信和重排
  类型：all_gather / ulysses / hybrid
  调度：sync / async

最终组合关系是：

ShardingPlanner
  -> ShardingPlan.cp_config
  -> ModuleShardingSpec 的 adapter/layout/handoff 信息
  -> ShardingApplier 解析 adapter + strategy
  -> 同一个 CP runtime 同时服务 production 和 validate

2. 当前实现与旧能力对比

能力	旧流程	新 components 当前实现	任务4目标
K/V all-gather	支持	支持	保持兼容
Pure Ulysses A2A	支持	不支持	迁移
Hybrid CP	支持	不支持	迁移
异步 K/V all-gather	支持	不支持	迁移
异步 Ulysses A2A	支持	不支持	迁移
Hybrid 异步	A2A 异步、K/V gather 同步	不支持	按旧语义迁移
Head-Tail 负载均衡	Colossal 支持	不支持	第二阶段迁移
DSA CP	同步/异步 all-gather	不支持	通过专用 adapter 迁移
TP + CP DTensor layout	支持	当前 all-gather 支持	所有 strategy 都要支持
production / validate 同核	旧流程不是新双模式结构	当前支持	新模式必须保持

当前新流程的关键问题不是缺少一个 all_to_all 函数，而是：

1. CP_WRAPPER_REGISTRY 同时承担了模型识别和通信行为，职责耦合；
2. _cp_sdpa_call 固定执行 K/V all-gather；
3. CP 只有一个一维 cp_mesh，没有 Hybrid 的子组拓扑；
4. spec 没有 CP 算法、同步方式、Q/K/V layout 和异步 handoff 描述；
5. planner 只校验序列长度，没有校验 Ulysses 的本地 head 数；
6. 当前 cp_utils.py 直接调用 torch.distributed，没有复用旧流程已有的 platform 可微和异步通信接口。

3. 设计目标

3.1 必须满足

• cp_size 单独启用时保持当前默认行为：同步 K/V all-gather；
• 一个 attention adapter 可以组合任意合法 CP strategy；
• production 和 validate 使用同一个 adapter、strategy 和本地计算路径；
• CP 内部通信不泄漏到模块边界 placement，边界仍保持 CP Shard(sequence)；
• TP、CP、FSDP、PP 组合时不丢失非 CP layout；
• 配置不合法时在 plan/apply 阶段 fail-fast，不能等到训练中静默算错；
• process group 在 apply 阶段创建或从 DeviceMesh 缓存读取，forward 中禁止 new_group；
• 同步、异步通信都必须有正确的 autograd 反向语义。

3.2 第一版不做

• 不在第一版支持不等长 sequence shard；
• 不自动猜测任意 HF 模型中“最后一个 Q/K/V producer”；
• 不承诺所有 GQA/MQA 都能运行 Ulysses；
• 不把 Hybrid 的内部 co/ulysses 维暴露成全局 ShardingPlan mesh 轴；
• 不把旧 ContextParallel 类直接嵌入新 applier，避免新旧控制流再次耦合。

4. 配置模型

4.1 Plan 级 CP 配置

新增 plan 级配置，算法选择不重复写入每个 attention spec：

class CPAlgorithm(str, Enum):
    ALL_GATHER = "all_gather"  # 原 Pure Colossal
    ULYSSES = "ulysses"
    HYBRID = "hybrid"


class CPSchedule(str, Enum):
    SYNC = "sync"
    ASYNC = "async"


@dataclass(frozen=True)
class ContextParallelConfig:
    algorithm: CPAlgorithm = CPAlgorithm.ALL_GATHER
    schedule: CPSchedule = CPSchedule.SYNC
    ulysses_degree: Optional[int] = None
    load_balance: bool = False
    async_fallback: Literal["error", "sync"] = "error"

ShardingPlan 新增：

@dataclass
class ShardingPlan:
    ...
    cp_config: Optional[ContextParallelConfig] = None

ShardingPlanner.plan 新增：

plan = planner.plan(
    model,
    mesh,
    tp_size=tp_size,
    cp_size=cp_size,
    cp_config=ContextParallelConfig(
        algorithm="hybrid",
        schedule="async",
        ulysses_degree=2,
    ),
)

兼容规则：

• cp_size == 1：忽略 CP runtime；
• cp_size > 1 且未传 cp_config：生成 all_gather + sync，保持当前新流程行为；
• 新 API 使用显式 algorithm，不再让 ulysses_degree 隐式决定模式；
• 兼容旧配置时可做一次转换：degree=1 -> all_gather，degree=cp -> ulysses，中间值 -> hybrid。

4.2 Module 级 attention 描述

inner_target 和 inner_wrapper 继续负责“目标位置”和“adapter 类型”。新增的字段只描述模块特有 layout 和异步 handoff：

@dataclass
class ModuleShardingSpec:
    ...
    inner_target: Optional[str] = None
    inner_wrapper: Optional[Union[str, Callable]] = None

    cp_seq_dim: int = 2
    cp_head_dim: int = 1
    cp_qkv_indices: Tuple[int, int, int] = (0, 1, 2)
    cp_qkv_kwarg_names: Tuple[str, ...] = ()
    cp_handoff_targets: Dict[str, str] = field(default_factory=dict)

    _resolved_inner_wrapper: Optional[str] = None
    _resolved_cp_strategy: Optional[str] = None

cp_handoff_targets 的路径相对 attention boundary：

ModuleShardingSpec(
    inner_target="self",
    inner_wrapper="sdpa_hf",
    cp_handoff_targets={
        "q": "q_norm",
        "k": "k_norm",
        "v": "v_handoff",
    },
)

handoff 必须是进入 attention primitive 前最后一个、输出直接作为 Q/K/V 使用的模块。若 projection 后还有 reshape、transpose、RoPE 或 QK-Norm，则不能把原始 q_proj/k_proj 当成安全 handoff。

5. 两级注册表

5.1 Attention Adapter 注册表

Adapter 只解决“怎样接入网络”，不决定 CP 算法：

CP_ADAPTER_REGISTRY = {
    "sdpa_qkv": SdpaQkvAdapter,
    "sdpa_hf": SdpaHfAdapter,
    "flex_qkv": FlexQkvAdapter,
    "flex_hf": FlexHfAdapter,
}

现有四个 wrapper 的迁移关系：

当前 wrapper	新 adapter 的职责
_wrap_sdpa_for_cp	获取 forward 的 Q/K/V，调用 runtime，再调用原 forward
_wrap_hf_sdpa_for_cp	拦截 F.scaled_dot_product_attention，把 primitive 调用交给 runtime
_wrap_flex_attn_for_cp	获取 Q/K/V，调用 runtime，再调用 FlexAttention
_wrap_hf_flex_for_cp	拦截 flex_attention，把 primitive 调用交给 runtime

HF adapter 继续保留 misfire detection。一次 forward 没有拦截到目标 primitive 时立即报错。

5.2 CP Strategy 注册表

CP_STRATEGY_REGISTRY = {
    "all_gather": AllGatherCPStrategy,
    "ulysses": UlyssesCPStrategy,
    "hybrid": HybridCPStrategy,
}

统一接口建议为：

class CPStrategy:
    def prepare_qkv(self, q, k, v, runtime_ctx):
        """返回 attention 实际消费的 q/k/v 和本次调用状态。"""

    def restore_output(self, output, call_state):
        """把 attention 输出恢复为本地 sequence shard。"""

    def install_async_handoffs(self, target, handoff_spec, runtime_ctx):
        """异步模式下安装 producer hooks；同步模式为空操作。"""

    def causal_metadata(self, call_state):
        """返回全局 Q offset、Q/KV 长度等 mask 信息。"""

Adapter 和 strategy 的组合伪代码：

def cp_aware_attention(q, k, v, primitive, kwargs):
    local_q, local_k, local_v, output_layout = unwrap_validate_dtensor(q, k, v)

    q2, k2, v2, state = strategy.prepare_qkv(
        local_q, local_k, local_v, runtime_ctx
    )
    kwargs = mask_adapter.rewrite(kwargs, strategy.causal_metadata(state))
    output = primitive(q2, k2, v2, **kwargs)
    output = strategy.restore_output(output, state)

    return rewrap_validate_output(output, output_layout)

这样增加新算法时只增加 strategy；增加新网络调用方式时只增加 adapter。

5.3 旧注册表兼容

迁移期间保留 CP_WRAPPER_REGISTRY：

• 内置 sdpa_qkv/sdpa_hf/flex_qkv/flex_hf 优先解析为新 adapter；
• 用户注册的旧式 full-wrapper 仍按 fn(target, cp_mesh) 调用；
• 旧式 full-wrapper 只允许 all_gather + sync，其他 strategy 必须报错，不能假装已经应用；
• inner_wrapper=Callable 继续表示用户完全接管，框架不再叠加内置 strategy。

6. 三种同步通信算法

以下统一以 BNSD [B, N, S, D] 为例；实际维度由 cp_seq_dim/cp_head_dim 指定。

6.1 AllGather / Pure Colossal

输入：

Q: [B, Nq,  S/P, D]
K: [B, Nkv, S/P, D]
V: [B, Nkv, S/P, D]

前向：

K/V --all_gather(sequence)--> [B, Nkv, S, D]
Q 保持本地 [B, Nq, S/P, D]
Attention(Q_local, K_global, V_global)
输出保持 [B, Nq, S/P, D]

反向：

dK_global/dV_global --reduce_scatter(sum)--> 本 rank 的 dK/dV sequence shard
dQ 不做 CP 通信

当前 _AllGatherAlongDim.backward 的 all_reduce + chunk 语义正确，但通信和显存开销较大。新 collective 层应优先使用 reduce_scatter_tensor，后端不支持时再回退到 all_reduce + chunk。

6.2 Pure Ulysses

约束：本 rank 在 TP 切分后的 Q head 数和 KV head 数都能被 cp_size 整除。

前向：

Q/K/V:
  [B, N, S/P, D]
    --split head + all_to_all-->
  [B, N/P, S, D]

Attention 在完整 sequence、局部 heads 上执行

Output:
  [B, N/P, S, D]
    --inverse all_to_all-->
  [B, N, S/P, D]

反向中每个 all-to-all 的梯度都是方向相反的 all-to-all，不需要 reduce。

建议低层提供两个清晰接口：

ulysses_seq_to_head(tensor, seq_dim, head_dim, group)
ulysses_head_to_seq(tensor, seq_dim, head_dim, group)

二者使用自定义 autograd Function 或 platform 的可微 all_to_all_single，互为 backward。

6.3 Hybrid

设：

P  = cp_size
U  = ulysses_degree
C  = P / U

从一维 CP mesh 派生内部二维拓扑：

hybrid_cp_mesh shape = (C, U)
mesh_dim_names       = ("cp_co", "cp_ulysses")

该二维拓扑只在 CP runtime 内使用，模块边界仍只认识全局 cp 轴。

前向：

1. Q/K/V 在 cp_ulysses 子组做 seq->head A2A
   [B, N, S/P, D] -> [B, N/U, S/C, D]

2. K/V 在 cp_co 子组沿 sequence all-gather
   [B, N/U, S/C, D] -> [B, N/U, S, D]

3. Q 保持 [B, N/U, S/C, D]，执行 attention

4. 输出在 cp_ulysses 子组做 head->seq inverse A2A
   -> [B, N, S/P, D]

反向顺序：

输出 inverse A2A 的 backward
  -> attention backward
  -> K/V all-gather 的 reduce-scatter
  -> Q/K/V seq->head A2A 的 inverse A2A

Hybrid group 必须在 apply 阶段从 mesh["cp"] 派生和缓存，forward 中只取 group handle。

7. 异步调度设计

7.1 可重叠位置

算法	forward 提前发起	consumer 等待	backward 提前发起
all_gather async	K/V handoff 后发起 all-gather	attention 前等待	all-gather backward 发起 reduce-scatter，K/V producer backward 前等待
Ulysses async	Q/K/V handoff 后发起 seq->head A2A	attention 前等待	A2A backward 发起 head->seq A2A，producer backward 前等待
Hybrid async	Q/K/V handoff 后发起 Ulysses A2A	attention 前等待；K/V co gather 同步	按旧实现：A2A 反向可异步，co reduce-scatter 同步

Hybrid 第一版保持旧流程的“半异步”语义。K/V 在 cp_co 上的 all-gather 暂不提前，避免两级 collective 的反向依赖和 stream 顺序变复杂。

7.2 Handoff 安全规则

异步模式不能仅凭属性名看到 q_proj 就自动启用。必须满足：

1. handoff 输出就是 attention adapter 最终看到的 Q/K/V；
2. handoff 与 consumer 之间没有 reshape、transpose、RoPE、norm 等未重放操作；
3. tensor layout 与 spec 的 cp_seq_dim/cp_head_dim 一致；
4. all-gather 需要 K/V handoff；Ulysses/Hybrid 需要 Q/K/V handoff。

apply 阶段处理规则：

handoff 完整且合法
  -> 安装 async producer/consumer hooks

handoff 缺失 + async_fallback="sync"
  -> WARNING + 回退同步
  -> _resolved_cp_strategy 写入实际 sync 策略

handoff 缺失 + async_fallback="error"
  -> fail-fast

新 API 默认 error，防止用户以为已经获得 overlap；旧 AsyncContextParallel 兼容层可使用 sync 回退，保持旧行为。

7.3 Runtime 状态

异步 handle 不能只放在一个覆盖式字典里。pipeline microbatch、activation checkpoint 重算和嵌套 forward 会同时存在多个调用。

每个已应用 attention 建立独立 CPRuntimeState：

@dataclass
class CPRuntimeState:
    forward_slots: Dict[str, Deque[AsyncSlot]]
    backward_slots: Dict[str, List[AsyncSlot]]
    invocation_id: int

• forward producer 按 FIFO 入队，attention consumer 按 FIFO 消费；
• backward 按 autograd 的逆序使用 stack；
• consumer 必须验证 slot 的 shape、dtype、device 和 invocation；
• forward 退出时有未消费 slot，视为 adapter/handoff 配置错误；
• wrapper 重复 apply 必须幂等或明确报错。

7.4 Autograd 语义

异步通信不能用“先异步裸 collective，稍后直接拿 output”的方式绕开 autograd。需要一个可微 wait 节点：

producer post-hook:
  launch raw async collective -> AsyncSlot(work, buffer)

attention pre-hook:
  DifferentiableWait.apply(original_tensor, AsyncSlot)
  -> wait
  -> 重排 output
  -> 把反向通信挂入 autograd graph

优先复用旧流程在 hyper_parallel.platform 已有的接口：

• all_gather_single(..., async_op=True)；
• all_to_all_single(..., async_op=True)；
• differentiable_async_allgather_wait；
• differentiable_async_a2a_wait；
• reduce_scatter_single。

不要在 components 中复制第二套 Torch-only async autograd 实现。

8. Causal mask

不同 strategy 的 Q 全局起点不同，mask 不能继续固定使用：

lo = cp_rank * q_local_len

正确规则：

算法	attention 内 Q 范围	causal offset
all_gather	当前 CP rank 的连续 chunk	cp_rank * q_len
Ulysses	完整 sequence	0，保留标准 is_causal=True
Hybrid	当前 cp_co rank 的连续大 chunk	co_rank * q_len_after_a2a
load-balance	每个 Q half 有独立全局区间	每次子 attention 单独计算

因此 strategy 应返回 CausalMetadata，adapter 只负责把它转换成 SDPA mask、Flex block mask 或后端 flash attention 参数。

若原调用已经有 attn_mask：

• bool mask：与 CP causal mask 做逻辑与；
• additive float mask：把 CP 不可见位置加上负无穷；
• 不能直接覆盖用户 mask；
• Ulysses 中 Q/KV 都是全长时，保持原 mask 和 is_causal；
• FlexAttention 的 block mask 必须按 strategy 执行后的 Q/KV 全局 shape 构造。

第一版优先保证正确性；显式 SDPA mask 可能让 flash kernel 降级。后续可由 NPU/Flash adapter 使用原生 offset 参数恢复性能。

9. TP、GQA、MQA 与 layout 校验

9.1 TP + Ulysses

Ulysses 的整除检查必须使用 TP 后的本地 head 数：

local_q_heads  = num_attention_heads / tp_size
local_kv_heads = num_key_value_heads / tp_size

local_q_heads  % ulysses_degree == 0
local_kv_heads % ulysses_degree == 0

等价的全局检查是 Q/KV head 数都能被 tp_size * ulysses_degree 整除。

9.2 GQA/MQA

第一版采用严格策略：

• Q 和 K/V 分别检查 head 整除；
• 任一不满足就拒绝 Ulysses/Hybrid；
• 提示用户改用 all-gather，或减小 ulysses_degree；
• 不在内部静默 repeat K/V heads。

后续如需支持 KV heads 少于 Ulysses degree，可增加显式 kv_head_policy="repeat"，并单独评估显存与通信开销。

9.3 Layout

内置 primitive adapter 的默认 layout：

SDPA / Flex primitive: BNSD, seq_dim=2, head_dim=1

QKV 风格自定义 core attention 可以通过 spec 改为 BSHD：

ModuleShardingSpec(
    inner_wrapper="sdpa_qkv",
    cp_seq_dim=1,
    cp_head_dim=2,
)

进入通信前必须验证：

• seq_dim != head_dim；
• Q/K/V rank 一致；
• batch、sequence 和 head_dim size 相容；
• Ulysses 目标 head 维可整除；
• Hybrid 的 cp_size % ulysses_degree == 0。

10. production / validate 双模式

所有 strategy 都必须遵循当前 D-01'' 原则：

production:
  local tensor -> explicit CP strategy -> local attention -> local output

validate:
  DTensor.to_local()
    -> 同一个 explicit CP strategy
    -> 同一个 local attention
    -> 按 spec.out_src 从 local 重新包装 DTensor

禁止 validate 依赖 DTensor redistribute、production 使用手写 collective 的双路径实现，否则只能验证 placement，不能验证 kernel 级等价。

TP layout 保存方式：

1. adapter 记录 Q/hidden_states 入场时的非 CP mesh 和 placements；
2. strategy 只处理 local tensor 和内部 CP group；
3. Ulysses/Hybrid 输出恢复为本地 sequence shard；
4. adapter 按原非 CP layout 或 spec out_src 重包装；
5. 内部 cp_co/cp_ulysses placement 不离开 attention 区域。

11. Planner 与 Applier 流程

11.1 Planner

def plan(..., cp_size=1, cp_config=None):
    cp_config = normalize_cp_config(cp_size, cp_config)
    validate_cp_config(model.config, tp_size, cp_size, cp_config)

    plan = infer_existing_sharding_plan(...)
    plan.cp_config = cp_config

    for attention_spec in plan:
        fill_default_cp_layout(attention_spec)
        merge_model_overrides(attention_spec)

    return plan

Planner 负责静态检查：

• algorithm、schedule、degree 组合；
• Hybrid degree；
• config 可见时的 Q/KV head 整除；
• load-balance 只允许 all-gather；
• CP 输入序列的 padding/切分约束。

11.2 Applier

def apply_phase_c(...):
    cp_mesh = get_cp_submesh(mesh)
    cp_topology = build_cp_topology(cp_mesh, plan.cp_config)

    for module, spec in attention_boundaries:
        adapter = resolve_cp_adapter(module, spec)
        strategy = resolve_cp_strategy(plan.cp_config)
        runtime = CPRuntime(
            topology=cp_topology,
            strategy=strategy,
            spec=spec,
            validate_mode=validate_mode,
        )
        adapter.install(module, runtime)

apply 阶段负责动态检查：

• inner target 与 adapter 是否匹配；
• async handoff 路径是否存在；
• 后端是否支持对应 sync/async collective；
• 模块实际 Q/K/V layout 是否与声明相符；
• 最终解析结果写回 _resolved_inner_wrapper/_resolved_cp_strategy 并记录日志。

12. 文件修改规划

文件	修改内容
sharding_config.py	增加 CPAlgorithm/CPSchedule/ContextParallelConfig；plan 级 cp_config；spec 的 layout/handoff 字段
sharding_planner.py	接收并归一化 cp_config；增加 mode/head/degree/load-balance 校验
cp_collectives.py（新增）	sync/async all-gather、reduce-scatter、seq-head A2A、可微 wait；统一走 platform
cp_strategy.py（新增）	AllGatherCPStrategy/UlyssesCPStrategy/HybridCPStrategy、内部 topology、call state
cp_adapters.py（新增）	四个现有 wrapper 迁成 adapter；primitive interception、DTensor unwrap/rewrap、misfire
cp_utils.py	保留 batch sharding和 mask 工具；flex_cp_allgather 做兼容导出
sharding_applier.py	只保留 adapter/strategy 解析和安装编排，移出具体 CP 算法
platform/platform.py	明确 CP 所需 collective capability；缺失时 fail-fast
components/distributed/__init__.py	对外导出 CP config 和注册表
tests/components/distributed/	增加各模式 forward/backward、双模式、TP×CP、异步状态测试

13. 兼容性与错误策略

13.1 默认行为

planner.plan(model, mesh, cp_size=2)

仍解析成：

algorithm=all_gather, schedule=sync

现有 CP 用例无需改配置即可继续运行。

13.2 必须 fail-fast 的情况

• cp_size % ulysses_degree != 0；
• Pure Ulysses 的 degree 不是 cp_size；
• Hybrid 的 degree 不在 (1, cp_size)；
• TP 后 Q 或 KV head 数不能被 degree 整除；
• load-balance 与 Ulysses/Hybrid 同时开启；
• async 缺少 handoff 且 fallback 为 error；
• adapter 未拦截到声明的 primitive；
• legacy full-wrapper 与非默认 CP strategy 组合；
• backend 不支持所需 collective；
• 同一 attention 被重复安装不同 CP runtime。

13.3 日志和可观测性

apply 后至少输出：

CP inner-wrap:
  module=model.layers.0.self_attn
  adapter=sdpa_hf
  strategy=hybrid
  schedule=async
  cp_size=8
  ulysses_degree=2
  colossal_degree=4
  handoff={q:q_norm,k:k_norm,v:v_handoff}
  fallback=none

实际回退同步时，日志和 _resolved_cp_strategy 必须反映真实执行方式。

14. 分阶段迁移

M1：配置和分层，不改变当前数值路径

• 引入 cp_config、adapter/strategy 两级抽象；
• 把当前 all-gather wrapper 迁到 AllGatherCPStrategy；
• 保持全部现有测试通过；
• 增加默认配置兼容测试。

M2：同步 Pure Ulysses

• 实现可微 seq-to-head / head-to-seq A2A；
• 接入四种内置 adapter；
• 增加 CP=2/4、causal/non-causal、forward/backward 测试；
• 增加 TP×CP head 整除测试。

M3：同步 Hybrid

• 从一维 CP mesh 派生并缓存 (cp_co, cp_ulysses)；
• 实现 A2A + K/V gather + inverse A2A；
• 增加 CP=4、U=2 的 forward/backward 和 mask offset 测试。

M4：异步模式

• 复用 platform async launch/wait；
• 实现 handoff 解析和 runtime slot 队列；
• 先支持 all-gather async、Ulysses async；
• Hybrid 按旧语义只异步 Ulysses A2A；
• 增加 microbatch、checkpoint recompute、forward/backward slot 生命周期测试。

M5：原有扩展能力

• Head-Tail load balance 作为 AllGather strategy 的 decorator；
• DSA Indexer/SparseAttention/Loss 作为专用 adapter 接入；
• 增加 NPU flash offset mask 优化；
• 补齐 MindSpore capability 或明确 noop/unsupported 行为。

15. 测试矩阵

15.1 单元测试

• config 默认值、旧 degree 转换、非法组合；
• adapter 自动识别和显式指定；
• CP topology 的 rank 分组；
• A2A 前后 shape 和 rank 数据顺序；
• all-gather backward 等于 reduce-scatter(sum)；
• A2A backward 等于 inverse A2A；
• async fallback 和 resolved strategy；
• causal metadata 的 offset。

15.2 多进程数值测试

world	策略	重点
CP=2	all-gather sync/async	forward、参数梯度、causal offset
CP=2	Ulysses sync/async	seq/head 重排、反向
CP=4, U=2	Hybrid sync/async	两级 group、mask、反向顺序
TP=2, CP=2	all-gather/Ulysses	TP layout 保存、本地 head 校验
CP=2	production vs validate	输出和梯度同核等价
CP=2	HF/qkv + SDPA/Flex	四个 adapter 均能组合 strategy

每个数值用例都与单卡完整序列参考比较：

forward local chunk
input gradient
Q/K/V/O projection parameter gradient
production vs validate

15.3 异步性能验证

数值正确只证明“异步接口能运行”，不能证明发生了 overlap。NPU 集成测试还应检查 profiler timeline：

projection GEMM
  与 CP all-gather / all-to-all 有时间交叠
attention 开始前对应 work 已 wait
backward projection GEMM 前 reduce-scatter / inverse A2A 已完成

16. 用户侧最终形态

同步 Ulysses：

cp_config = ContextParallelConfig(
    algorithm="ulysses",
    schedule="sync",
)
plan = planner.plan(model, mesh, cp_size=4, cp_config=cp_config)

异步 all-gather：

cp_config = ContextParallelConfig(
    algorithm="all_gather",
    schedule="async",
)

Hybrid：

cp_config = ContextParallelConfig(
    algorithm="hybrid",
    schedule="async",
    ulysses_degree=2,
)

网络特殊信息仍通过 plan override 描述：

overrides = {
    "model.layers.0.self_attn": ModuleShardingSpec(
        inner_target="self",
        inner_wrapper="sdpa_hf",
        cp_seq_dim=2,
        cp_head_dim=1,
        cp_handoff_targets={
            "q": "q_norm",
            "k": "k_norm",
            "v": "v_handoff",
        },
    ),
}

算法配置属于 plan，网络结构差异属于 module spec。两者分离后，同一个模型适配不需要为每种 CP 通信模式重新写 wrapper。

17. 验收标准

任务4后续实现完成的最低标准：

1. cp_size>1 默认路径与当前 all-gather 行为完全兼容；
2. all-gather、Ulysses、Hybrid 三种同步模式前向和反向对齐单卡；
3. all-gather、Ulysses 异步模式数值对齐，Hybrid 按旧半异步语义对齐；
4. 四个内置 attention adapter 不复制三套通信代码；
5. production 与 validate 对每种模式输出、梯度一致；
6. TP×CP 下 layout 和本地 head 数正确；
7. 所有非法 mode/head/handoff/backend 组合在训练前失败；
8. profiler 能证明异步模式存在实际计算通信重叠。