Hyper-Parallel Chunk KDA 当前实现设计

1. 文档目的与版本基线

本文档描述当前已经落地的 Chunk KDA 实现，而不是早期接入设想。重点说明：

1. FLA-NPU、Hyper-Parallel（HP）和 MindFormers（MF）的职责边界。
2. HP 对外暴露的总接口和九个显式阶段接口。
3. 单卡训练、Ulysses CP 和 sequence-split state-scan CP 下的数据流转。
4. 当前能力边界、已验收范围和未解决问题。

组件	基线
FLA-NPU	newest_flash_linear_attention / 70eb6ba8
Hyper-Parallel	hyper_parallel_github / kda_r1.0 / a26a38dd
MindFormers	mf 目录中的 KDA Ulysses 和实验性 state-scan 调用方
后端	MindSpore Ascend；PyTorch 后端明确不支持

---

2. 总体设计

HP 对外提供两层接口：

                         模型 / MindFormers
                                |
              +-----------------+------------------+
              |                                    |
              v                                    v
    npu_chunk_kda 可微总接口              九个显式 stage 接口
      HP 管理缓存和自动微分                调用方管理通信和反向
              |                                    |
       +------+-------+                     FLA stage ACLNN
       |              |                            |
       v              v                            v
 FwdV2+fused bwd   六阶段 autograd              AscendC kernels
       |              |
       +------+-------+
              v
         FLA-NPU ACLNN/OPP

这种分层的原因是：

• 普通训练和 Ulysses CP 需要一个易用的、可自动求导的总接口。
• 沿 sequence 切分的 CP 必须在 KDA 状态计算中间插入 AllGather 或 PTP，单次总算子调用无法插入这些通信，因此需要显式 stage。
• HP 负责通用计算能力，不应绑定通信组、rank 拓扑或某一种 CP 算法。

---

3. 组件职责划分

3.1 FLA-NPU

FLA-NPU 是 KDA 数学语义和底层算子能力的提供者，负责：

• Chunk KDA 前向、反向和分阶段内核。
• ACLNN ABI、shape/dtype 校验、tiling 和 AscendC 实现。
• gate、Q/K L2Norm、beta sigmoid、state 递推等数学语义。
• 计算本地 affine summary，以及对已经 gather 的 summary 执行 merge。

FLA-NPU 不负责 HP 公共 Python API，也不负责 MF 的并行通信策略。

3.2 Hyper-Parallel 公共 API 层

位置：

• hyper_parallel/custom_ops/experimental/experimental_ops.py
• hyper_parallel/custom_ops/experimental/chunk_kda_stages.py
• hyper_parallel/custom_ops/experimental/__init__.py

职责：

• 对外暴露稳定的 Python 函数名和参数语义。
• 提供统一可微总接口 npu_chunk_kda。
• 提供九个调用方可编排的显式 stage 原语。
• 在进入平台层前完成公共参数约束校验。

该层不实现 kernel，不创建通信组，不执行 collective。

3.3 HP MindSpore 路由与自动微分层

位置：

• 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/chunk_kda_autograd.py

职责：

• 根据参数选择 fused 路径或 stage-chain 路径。
• 完成 BSND/BNSD/TND/NTD 的规范化与梯度布局恢复。
• GQA 时展开 Q/K heads，反向时对重复 heads 求和。
• 保存或重算反向需要的 cache。
• 将 final_state 的上游梯度传入 Dhu，并将 dh0 返回给 initial_state。
• 在总接口中隐藏 stage 中间量，只向模型返回 output 和可选 final state。

当前路由规则：

safe_gate=True
and use_gate_in_kernel=True
and use_exp2=True
and a_log is not None
             |
             +--> 使用六阶段 stage-chain autograd

或者 initial_state 非空 / 开启重算 / state_v_first=True
             |
             +--> 使用六阶段 stage-chain autograd

其他受支持组合
             |
             +--> FwdV2 + fused backward

safe raw-gate 在无 state 时也选择 stage chain，是因为该路径与当前 FLA stage ABI 成对，并且规避了 MF Ulysses 工作负载中 fused backward 出现的非有限梯度。

3.4 HP C++ Bridge 层

位置：

• chunk_kda_fwd.cc
• chunk_kda_bwd.cc
• chunk_kda_stages.cc
• chunk_kda_stage_dispatch.h

职责：

• 将 MindSpore Tensor 转为 ACLNN 可消费的输入。
• 按 ABI 分配固定顺序的输出 Tensor。
• 调用 aclnnChunkKdaFwdV2、aclnnChunkKdaBwd 和九个 stage ACLNN。
• stage 符号按需动态解析；某个符号缺失时只让对应调用失败。

Bridge 是薄适配层，不重新实现 KDA 数学逻辑，也不隐式执行通信。

3.5 HP ContextParallel

HP 现有 ContextParallel 负责 Ulysses 的 sequence/head all-to-all：

• KDA 调用前：local sequence/full heads 转为 full sequence/local heads。
• KDA 调用后：full sequence/local heads 转回 local sequence/full heads。

npu_chunk_kda 本身不知道 CP group，也不知道 all-to-all。

3.6 MindFormers 调用方

MF 负责模型级和并行级编排：

• 将 HP ContextParallel 挂到 KDA core 外层。
• 设置 kda_disable_recompute。
• sequence-split CP 中创建 CP group，执行 summary AllGather/Merge 或未来的 PTP scan。
• 保证 rank 顺序与 sequence 时间顺序一致。
• 负责 a_log/dt_bias 等 replicated 参数梯度归约。
• 负责 ShortConv halo、loss weighting、packed sequence ownership 等模型语义。

---

4. 总接口定义

npu_chunk_kda(
    q, k, v, g, beta,
    scale: float,
    chunk_size: int = 64,
    *,
    layout: str = "BNSD",
    output_final_state: bool = False,
    safe_gate: bool = False,
    lower_bound: float = -5.0,
    use_gate_in_kernel: bool = True,
    a_log=None,
    dt_bias=None,
    cu_seqlens=None,
    chunk_indices=None,
    initial_state=None,
    disable_recompute: bool = True,
    state_v_first: bool = False,
    epsilon: float = 1e-6,
    use_qk_l2norm_in_kernel: bool = False,
    use_beta_sigmoid_in_kernel: bool = False,
    allow_neg_eigval: bool = False,
    use_exp2: bool = True,
) -> tuple[Tensor, Optional[Tensor]]

4.1 输入含义

参数	含义
q/k	Query/Key，head 数为 Hk，末维为 K
v	Value，head 数为 Hv，末维为 V
g	raw gate 或 gate 输入；当前可训反向要求 kernel 内处理 raw gate
beta	KDA 更新系数；可选在 kernel 内执行 sigmoid
scale	QK 缩放，通常为 K**-0.5
layout	BNSD/BSND/NTD/TND
a_log/dt_bias	raw gate 变换的 FP32 参数；dt_bias 可为 [Hv,K] 或等价展平形式
cu_seqlens/chunk_indices	packed/variable-length 的 host metadata
initial_state	每个 batch 或 packed sequence 的 KDA 边界状态
disable_recompute	True 保存反向 cache；False 在反向重建 prepare 和 state relay
state_v_first	对外 state 最后两维使用 [V,K]；HP 内部规范化为 K-first
epsilon	Q/K L2Norm 的数值稳定项
use_qk_l2norm_in_kernel	是否在 FwdV2/stage prepare 中完成 Q/K 归一化
use_beta_sigmoid_in_kernel	是否在 kernel 内对 beta 执行 sigmoid
allow_neg_eigval	负特征值配套模式；必须与 beta sigmoid 一起使用
use_exp2	使用 exp2 语义；当前 backward 要求为 True

4.2 Layout

layout	Q/K	V/G	beta
BNSD	[B,Hk,T,K]	[B,Hv,T,V/K]	[B,Hv,T]
BSND	[B,T,Hk,K]	[B,T,Hv,V/K]	[B,T,Hv]
NTD	[Hk,T,K]	[Hv,T,V/K]	[Hv,T]
TND	[T,Hk,K]	[T,Hv,V/K]	[T,Hv]

Hv 可大于 Hk。stage-chain 自动微分会将 Q/K 复制到 value heads，反向时再将重复 head 梯度求和回原始 Hk。

4.3 返回值

(attn_out, final_state)

• attn_out 为 token-major 业务输出，HP 在 stage chain 外恢复公共输出约定。
• output_final_state=False 时，final_state=None。
• stage-chain 路径中，final_state 可参与 loss，其梯度会传给 Dhu。
• 如果提供 initial_state，stage-chain backward 返回对应 dh0。
• fused backward 路径仍不允许 final_state 参与 loss。

4.4 当前训练交集

当前 stage-chain 状态/重算路径的已对齐配套组合为：

• BF16 Q/K/V。
• K=V=128。
• chunk_size=64。
• safe_gate=True。
• use_gate_in_kernel=True。
• use_exp2=True。
• 必须提供 a_log；BwdFinalize 配套 ABI 要求 dt_bias，未传时 HP 构造零值。

FLA 前向能力比这个交集更广，但不能由“前向可运行”推导“反向已验收”。

---

5. 显式 Stage 接口

Stage API 都是本地 Tensor 算子，没有隐式自动微分，也没有通信。

5.1 前向阶段

npu_chunk_kda_fwd_prepare

npu_chunk_kda_fwd_prepare(
    q, k, v, g, beta, scale, chunk_size=64, *,
    layout="BNSD", a_log=None, dt_bias=None,
    cu_seqlens=None, chunk_indices=None,
    epsilon=1e-6,
    use_qk_l2norm_in_kernel=False,
    use_gate_in_kernel=False,
    use_beta_sigmoid_in_kernel=False,
    allow_neg_eigval=False,
    safe_gate=False, lower_bound=-5.0,
    use_exp2=True, backward_mode="save",
) -> tuple[13]

作用：将本地 Q/K/V/g/beta 处理为 head-major cache。

固定返回顺序：

gk, Aqk, Akk, w, u, qg, kg, qg_scaled,
q_hat, k_hat, q_rstd, k_rstd, beta_eff

backward_mode 只控制输出哪些 cache，不代表启用自动微分：

模式	非 None 槽位
none	0, 1, 3, 4, 6, 7
forward	none + 2
recompute	none + 8..12
save	全部 13 项

npu_chunk_kda_pre_process_fwd

npu_chunk_kda_pre_process_fwd(
    k, w, u, *, gk, cu_seqlens, chunk_size=64
) -> tuple[Tensor]

返回 FP32 summary [N,Hv,K,V+K]，其中最后一维表示 [E | M]：

H_out = M @ H_in + E

该 summary 用于 sequence-split CP 的跨 rank 状态扫描。

npu_chunk_kda_fwd_h

npu_chunk_kda_fwd_h(
    k, w, u, *, gk,
    initial_state=None,
    output_final_state=False,
    chunk_size=64,
    cu_seqlens=None, chunk_indices=None,
    use_exp2=True, state_v_first=False,
) -> tuple[Tensor, Tensor, Optional[Tensor]]

作用：从调用方给定的边界 initial_state 开始，完成本地 chunk 状态递推。

返回：

• h: [B,Nc,Hv,K,V]，每个 chunk 的入口/递推状态。
• v_new: 更新后的 value relay。
• final_state: 本地序列终点状态。

npu_chunk_kda_fwd_finalize

npu_chunk_kda_fwd_finalize(
    qg_scaled, aqk, v_new, h, *,
    output_layout="BSND",
    state_v_first=False,
    cu_seqlens=None, chunk_indices=None,
) -> Tensor

作用：将 Q 侧 cache、局部 value 和 state relay 组合为最终 attention output。

5.2 反向阶段

npu_chunk_kda_bwd_prepare

npu_chunk_kda_bwd_prepare(
    aqk, v_new, d_o, h, *,
    scale, layout="BNSD", chunk_size=64,
    state_v_first=False,
    cu_seqlens=None, chunk_indices=None,
) -> tuple[Tensor, Tensor, Tensor]

返回：

(dAqk, dv_local, dq_raw)

该阶段先拆出与 output 直接相关的本地梯度，为后续状态反向和最终梯度汇总做准备。

npu_chunk_kda_pre_process_bwd

npu_chunk_kda_pre_process_bwd(
    q, k, w, d_o, dv, scale, chunk_size=64, *,
    gk, cu_seqlens=None,
) -> Tensor

返回本地 backward affine summary [Hv,K,V+K]。调用方按反时间顺序交换 summary，为每个 rank 构造终点状态梯度 dht。

npu_chunk_kda_bwd_dhu

npu_chunk_kda_bwd_dhu(
    q, k, w, d_o, dv, scale, chunk_size=64, *,
    gk, h0=None, dht=None,
    cu_seqlens=None, chunk_indices=None,
    use_exp2=True, state_v_first=False,
) -> tuple[Tensor, Optional[Tensor], Tensor]

返回：

(dh, dh0, dv_scan)

• dht 是本 rank 终点边界梯度。
• dh 是 chunk 状态梯度，始终使用 K-first。
• 传入 h0 时才返回 dh0。
• dv_scan 是经状态反向累积后的 value 梯度。

npu_chunk_kda_bwd_finalize

npu_chunk_kda_bwd_finalize(
    q, k, v, gk, raw_g, beta,
    a_log, dt_bias, akk, v_new, h, dh,
    dv_scan, d_aqk, dq_raw, *,
    scale, lower_bound=-5.0, chunk_size=64,
    safe_gate=True, use_gate_in_kernel=True,
    use_exp2=True, state_v_first=False,
    q_rstd=None, k_rstd=None,
    cu_seqlens=None, chunk_indices=None,
    use_beta_sigmoid_in_kernel=False,
    allow_neg_eigval=False,
    output_layout="BNSD",
) -> tuple[7]

返回：

(dq, dk, dv, dbeta, dg, dA_log, ddt_bias)

该阶段汇总块内梯度、state-scan 梯度和 gate/beta 链式求导，得到面向模型输入和参数的最终梯度。

5.3 Summary 合并阶段

npu_chunk_kda_merge_fwd_bwd

npu_chunk_kda_merge_fwd_bwd(
    state, summaries,
    pre_or_post_num_ranks: int,
    rank: int,
    *, forward: bool = True,
    state_v_first: bool = False,
) -> Tensor

作用：对调用方已经 AllGather 得到的 summaries 进行 affine compose，原位更新 state。

该接口不执行 AllGather；rank、参与合并的数量和 summary 顺序均由调用方保证。

---

6. 不同场景的数据流

6.1 单卡/普通训练

调用方只使用 npu_chunk_kda，不需要直接处理 stage cache。

q,k,v,g,beta,a_log,dt_bias
              |
              v
       npu_chunk_kda
              |
       HP 选择执行路径
        /                 \
       v                   v
 FwdV2 + fused bwd    FwdPrepare -> FwdH -> FwdFinalize
                              |
                              v
                       attn_out, final_state
                              |
                            loss
                              |
                              v
                    BwdPrepare -> Dhu -> BwdFinalize
                              |
                              v
             dq,dk,dv,dg,dbeta,dA_log,ddt_bias,dh0

保存模式：

forward 保存 cache + h + v_new
                  |
                  v
backward 直接消费保存量

重算模式：

forward 只保存原始输入和边界状态
                  |
                  v
backward 重新执行 FwdPrepare + FwdH
                  |
                  v
      BwdPrepare -> Dhu -> BwdFinalize

6.2 Ulysses Context Parallel

Ulysses 不拆 KDA 的状态递推。它通过 all-to-all 让每个 rank 拿到完整 sequence 和部分 heads，然后直接调用总接口。

每 rank: local sequence, full heads
                  |
                  | ContextParallel sequence->head all-to-all
                  v
每 rank: full sequence, local heads
                  |
                  v
           npu_chunk_kda
                  |
                  v
   full sequence, local output heads
                  |
                  | ContextParallel head->sequence all-to-all
                  v
   local sequence, full output heads
                  |
                  v
                 loss

反向数据流由 MindSpore autograd 经 HP DFunction 和 ContextParallel 的反向规则自动逆转。

职责边界：

• ContextParallel 负责 all-to-all。
• npu_chunk_kda 只计算 full-sequence/local-head KDA。
• MF 负责 head 整除关系、参数分片和 ShortConv halo 等模型级问题。

6.3 Sequence-split state-scan CP：前向

每个 rank 保留一段按时间顺序排列的 token。KDA 状态需要从前一 rank 的终点传到后一 rank。

每 rank 的 local q/k/v/g/beta
                  |
                  v
             FwdPrepare
                  |
        +---------+----------+
        |                    |
        v                    v
 local caches        PreProcessFwd
                             |
                             v
                 local summary [E|M]
                             |
                             | MF AllGather
                             v
                    all-rank summaries
                             |
                             | Merge / FP32 compose
                             v
                   本 rank 边界 h0
        +--------------------+
        |
        v
      FwdH(h0) -> h, v_new, local final_state
        |
        v
   FwdFinalize -> local output

summary 表示本 rank 的 affine state map：

H_out = M @ H_in + E

因此可先交换较小的 summary，再计算每个 rank 的边界 h0，而不需要 gather 全部 token cache。

6.4 Sequence-split state-scan CP：反向

每 rank 收到 local d_o 和可选 final-state gradient
                  |
                  v
             BwdPrepare
                  |
        dAqk, dv_local, dq_raw
                  |
       +----------+-----------+
       |                      |
       v                      v
 local caches        PreProcessBwd
                              |
                              v
                    local backward summary
                              |
                              | MF AllGather
                              v
                    all-rank summaries
                              |
                              | 反向 Merge / compose
                              v
                    本 rank 终点 dht
       +----------------------+
       |
       v
 BwdDhu(h0,dht) -> dh, dh0, dv_scan
       |
       v
 BwdFinalize
       |
       v
dq,dk,dv,dg,dbeta,dA_log,ddt_bias
       |
       +--> MF 对 replicated 参数梯度执行 all-reduce
       +--> 只由 sequence 首 rank 贡献全局 initial_state 梯度，再归约

前向 summary 按 rank 0 到 rank N-1 的时间顺序组合；反向 summary 必须按相反顺序组合。

6.5 Packed/Variable-length

单卡 stage 通过 cu_seqlens/chunk_indices 识别序列边界：

packed tokens + cu_seqlens
          |
          v
 FwdPrepare / FwdH / Finalize
          |
    每条序列独立 reset state

但当前 MF sequence-split CP 未实现跨 rank packed ownership/reset 计划，所以不允许将“单卡 stage 支持 packed”等同于“多卡 state-scan CP 已支持 packed”。

---

7. 为什么不把 CP 通信放进 KDA 算子

1. 算子与拓扑解耦：KDA kernel 不需要感知 world size、rank 或 process group。
2. 支持多种编排：上层可以选择 AllGather+Merge、PTP scan 或其他 summary compose 方式。
3. 可测试性：每个 stage 可以在 CP=1 下独立与总接口和 CPU recurrence 对比。
4. 可选依赖：缺少某个 stage OPP 只影响对应路径，不阻断其他 HP custom ops。
5. 模型语义属于调用方：ShortConv halo、packed reset、loss scaling 和参数梯度归约都不是 KDA 本地 kernel 的职责。

---

8. 当前验证状态

8.1 已验证

• HP API/stage/extension UT：41 passed，另有 4 个 subtests passed。
• MF 配置 UT：21 passed。
• 单卡 NPU ST：41 passed。
• Ulysses CP2/CP4。
• TP2×CP2 + FSDP + LoRA + bias。
• Ulysses CP2/CP4 重算两步训练。
• state-scan CP2/tail65/保存模式。
• state-scan CP4/tail65/重算+V-first。
• state-scan CP2/aligned64/重算+V-first。

8.2 已知未通过项

强 gate 衰减、BF16、T=65、常量 dht、loss 仅依赖 final_state 时：

dk relative L2 = 6.8738% > 4% 验收门槛

当前定位为 BF16 v_new relay 误差经终态梯度点积的正负抵消放大。该失败用例保留，不 skip，不放宽阈值。因此目前不能宣称任意 final-state-only loss 都已完成精度验收。

8.3 实验性 state-scan CP 边界

当前 MF KdaStateScan 仅支持：

• B=1。
• 每 rank 一段非空、按时间顺序排列的 dense sequence。
• BF16、K=V=128、equal heads、safe raw gate。
• AllGather + Merge 编排。

尚未完成：

• causal ShortConv 跨 rank halo。
• packed sequence 的跨 rank ownership/reset。
• PTP 状态扫描。
• 异步流水和性能收益验收。

---

9. 结论

当前 HP Chunk KDA 的实现原则是：

使用一个可微总接口承载普通训练和 Ulysses CP，使用一组无通信的显式 stage 原语为 sequence-split CP 提供可编排能力。

FLA-NPU 负责 kernel 和数学语义，HP 负责 API、layout、自动微分与本地 stage bridge，MF 负责模型与通信编排。这样既能保持普通训练的易用性，又不会将 HP 绑定到某一种 CP 通信算法。