KDA 融合算子与 Context Parallel PR 测试报告

KDA 融合算子与 Context Parallel PR 测试报告

文档结构总览

本报告面向 KDA Context Parallel PR 的功能验收。各章节直接给出测试对象、配置、方法、
结果和结论，仅覆盖本 PR 提供的功能。

章节	验证对象	主要目的
第 1 章	PR 功能范围	明确本 PR 支持的 CP 模式、backend 和验收边界
第 2 章	环境与统一方法	固定软硬件、精度指标、计时及显存口径
第 3 章	Backend 与交付检查	验证 Eager/Triton 分层、可选依赖、fail-fast 和 CPU UT
第 4 章	单卡 KDA	验证 Eager 数学参考和 Triton 融合算子的前反向精度
第 5 章	CP 完整 KDA 层	验证 Ulysses、P2P 的 output、hidden gradient 和参数梯度
第 6 章	多步训练	验证非零学习率下的 AdamW loss、grad norm 和参数轨迹
第 7 章	性能与显存	比较正式 KDA shape 下 P2P 与 Ulysses 的性能和峰值显存
第 8 章	Activation Checkpoint	验证 P2P 在 forward 重放时的通信和梯度正确性
第 9 章	PR 结论	汇总可用于 Issue 和 PR 描述的验收结论

1. 测试目标与范围

项目	内容
关联 PR	Hyper-Parallel PR #1311
被测提交	a867e401a287b6027d535c353c355ec85d72cd12
CP 模式	ulysses、p2p
KDA backend	eager、triton
训练范围	fixed-length dense KDA 前向与反向

本报告需要证明：

1. Hyper-Parallel 内置 Eager KDA 与逐 token recurrence、FLA naive reference 在 output、
   final state 和全部输入梯度上保持一致。
2. Triton KDA 的 output 和全部输入梯度与 FP32 Eager reference 在既定 BF16 误差范围内一致。
3. Ulysses 与 P2P 能正确执行完整 KimiDeltaAttention 层的前向和反向。
4. P2P 的 forward state chain、backward state-gradient chain 和 affine state summary 正确。
5. 非零学习率 AdamW 多步训练中，loss、grad norm 和参数误差不持续扩大。
6. Activation Checkpoint 重放 P2P forward 时不发生消息错配、死锁或梯度丢失。
7. 生产 shape 下 P2P 相对 Ulysses 的性能和显存结果可稳定复现。
8. 缺少 Triton 运行时、FLA 接口不兼容或 shape 不支持时明确报错，不静默回退 Eager。

本 PR 的功能矩阵如下：

CP mode	Eager backend	Triton backend	本报告验证内容
Ulysses	支持	支持	数学精度、完整层、多步训练、Triton 性能与显存
P2P	支持	支持	state summary、完整层、多步训练、checkpoint、Triton 性能与显存

以下内容不属于本 PR 的验收范围：

• packed/varlen KDA；
• 推理 KV/state cache；
• 完整 Kimi K3 模型、MLA、MoE 和多模态外层；
• AllGather KDA CP；
• backend 自动回退。

2. 环境与统一测试方法

2.1 软件与硬件环境

项目	配置
NPU	8 x Ascend 910B3，64 GiB HBM/device
CANN	9.1.0-beta.3
Python	3.11
PyTorch	2.10.0+cpu
torch-npu	2.10.0
Triton Python module	3.2.0，Ascend backend available
triton-ascend distribution	3.2.1+git2badfc89
flash-linear-attention	35dceaee5408e69a555fec34cb215c93c375dabe
token compute dtype	BF16
recurrent state / gate accumulation	FP32

Eager 测试不要求安装 FLA 或 Triton。Triton 测试固定使用上述 FLA revision，避免不同版本的
staged KDA 接口和 kernel 行为影响结果。

2.2 精度方法与判据

张量误差统一记录：

max_abs     = max(abs(actual - reference))
relative-L2 = ||actual - reference||_2 / max(||reference||_2, eps)
finite      = actual/reference 均无 NaN 和 Inf

单卡和 CP 测试使用相同初始权重、相同输入及相同随机上游梯度。CP 参数梯度在 CP group
内求和后再与单卡完整序列 reference 比较。

验证层级	通过条件
单卡 Triton KDA	output relative-L2 <= 2e-2；任一输入/参数梯度 relative-L2 <= 5e-2
CP 完整 KDA 层	output、hidden gradient、任一参数梯度 relative-L2 均 <= 3e-2
20 步 AdamW	最大 loss rel <= 1e-4；最大 grad-norm rel <= 2e-3；最终参数 rel-L2 <= 3e-2
所有测试	output、state、loss 和 gradient 全部 finite，不允许漏梯度

单独比较 loss 或 grad norm 不能证明 CP 数学等价，因此完整层测试还直接比较 output、hidden
gradient 和逐参数 gradient；多步 loss/grad norm 只用于检查训练轨迹稳定性。

2.3 性能与显存口径

性能和显存测试统一采用：

warmup = 3
repeat = 10
计时前后执行 device synchronize
分布式 elapsed 取所有 rank 的 MAX
编译与 warmup 时间不计入结果

显存记录所有参与 rank 的最大 max_memory_allocated 和 max_memory_reserved。Forward peak
和完整 Fwd+Bwd peak 分开记录，不使用 reserved 代替 allocated 判断真实活跃张量规模。

3. Backend、兼容性与交付检查

3.1 Backend 边界

eager 是 Hyper-Parallel 内置的 PyTorch 数学实现，承担以下职责：

• 单卡和 CPU 数学 reference；
• 支持随机非零 initial state 和 final-state gradient；
• 支持非 64 整除序列，并在内部完成 chunk padding；
• 不导入 FLA 或 Triton。

triton 是生产融合路径：

• Local KDA 和 Ulysses 调用 FLA 的公开 chunk_kda API；
• P2P 调用 FLA 的 staged KDA 前反向算子；
• P2P affine state summary 与 gradient summary 由 Hyper-Parallel 实现；
• FLA 仅在首次选择 backend="triton" 时延迟导入；
• 任一依赖或 capability 不满足时直接报错，不回退 Eager。

3.2 Triton capability 与 fail-fast

条件	预期行为
未安装 FLA	RuntimeError，明确提示缺少可选依赖
FLA 版本低于 0.6.0	RuntimeError，明确提示最低版本
Triton-Ascend backend 不可用	RuntimeError
缺少 P2P 所需 staged KDA API	RuntimeError，指出缺失模块或符号
q/k/v/gate/beta 不是 BF16	TypeError
A_log/dt_bias 不是 FP32	TypeError
Dk != 128、Dv != 128 或 chunk_size != 64	NotImplementedError
local sequence 不能被 64 整除	ValueError
tensor 不在同一 NPU	RuntimeError 或 ValueError
对同一模块重复 apply KDA CP	RuntimeError，避免嵌套两层 CP wrapper

3.3 CPU UT、静态检查与 wheel

检查项	结果
KDA Eager、state summary 和 CP apply UT	PASS，10 cases
FLA adapter 契约 UT	PASS，3 cases
CPU UT 合计	PASS，13 cases
git diff --check	PASS
Python compileall	PASS
pylint	PASS
Lizard 白名单审计	PASS
wheel 构建及隔离安装	PASS

CPU UT 重点覆盖：

1. 完整 Eager KimiDeltaAttention 前向、反向和全部参数梯度。
2. Eager chunkwise KDA 与 token recurrence 的 output、final state 和全部梯度。
3. 随机非零 initial_state 与随机非零 dHT。
4. CP4 local summary 的顺序组合、final state、token output 和全部梯度。
5. autocast 下 affine state 保持 FP32。
6. 非 chunk 整除序列的 Eager padding。
7. FLA 缺失、版本不兼容时的明确错误。
8. Hyper KDA 参数到 FLA public API 的完整转发。
9. 对同一模块重复 apply CP 时拒绝二次包装。

wheel 中包含 Hyper-Parallel 的 KDA layer、CP wrapper、FLA adapter、state-summary 实现和
KDA README；不复制 FLA 的 KDA kernel。未安装 FLA 时，导入 Hyper-Parallel 和运行 Eager
KDA 均不受影响。

4. 单卡 KDA 验证

4.1 Eager 数学参考

Eager KDA 采用两层交叉验证：

1. torch_recurrent_kda 与 FLA naive_recurrent_kda 比较逐 token recurrence。
2. torch_chunk_kda 与 FLA naive_chunk_kda 比较 chunkwise UT/WY 分解。

配置为 B=2,T=16,Hq=2,Hv=4,Dk=4,Dv=3,chunk=8,FP32，包含随机非零 h0、
随机非零 dHT，并比较 q/k/v/gate/beta/A_log/dt_bias/h0 的全部梯度。

对比	output max abs	final state max abs	全部梯度最大 abs	结果
Hyper recurrent vs FLA recurrent	5.960e-8	5.960e-8	1.192e-7	PASS
Hyper chunk vs FLA chunk	1.192e-7	5.960e-8	3.772e-7	PASS

另以 CP4 方式将 T=64 切为四段，顺序应用四个 local affine summary。拼接 token output、
逐段 final state、组合 summary 后的 final state 及全部梯度均与完整序列 Eager reference 一致：

检查项	容差	结果
token output / final state	atol=3e-5, rtol=3e-5	PASS
q/k/v/gate/beta/A_log/dt_bias/h0 gradients	atol=4e-4, rtol=4e-4	PASS

这些结果同时验证了 state summary 的正向组合和反向梯度链，没有遗漏 initial-state gradient
或 local summary 输入梯度。

4.2 Triton 融合算子精度

配置为 B=1,T=128,H=4,Dk=Dv=128,chunk=64。Triton 使用 BF16 token 输入与 FP32
A_log/dt_bias，reference 为 FP32 Eager recurrence；两边使用相同随机 output gradient。

指标	output	dq	dk	dv	dgate	dbeta	dA_log	ddt_bias
relative-L2	5.29e-3	4.69e-3	4.90e-3	5.08e-3	5.33e-3	5.26e-3	1.27e-2	5.61e-3

全部指标低于单卡门槛，且 output 与全部梯度 finite，结果 PASS。

公开单卡 Triton API 使用零 initial state 并返回 token output；P2P 所需的非零 rank initial
state、final state 以及反向 state gradient 由第 5 章的 stateful CP 路径验证。

5. CP 完整 KDA 层验证

5.1 测试范围

被测完整层包含：

input hidden states
  -> Q/K/V/Beta/Gate projection
  -> Q/K/V causal ShortConv
  -> KDA core
  -> output gate
  -> gated RMSNorm
  -> output projection
  -> attention output

单卡 Reference 使用官方 Hugging Face KimiDeltaAttention，与 CP 使用相同初始权重、输入
和随机上游梯度。Reference 运行完整序列；CP 将序列切为 local shard，参数梯度在 CP group
内求和后逐参数比较。

5.2 Eager backend

配置为 B=1,global T=1024,hidden=512,H=4,Dk=Dv=128,conv=4,chunk=64,BF16,CP4。

mode	output rel-L2	hidden grad rel-L2	最大参数 grad rel-L2	结果
Ulysses + Eager	8.905e-3	1.034e-2	1.039e-2	PASS
P2P + Eager	8.904e-3	1.055e-2	1.282e-2	PASS

两条 Eager CP 路径均低于 3e-2 门槛；所有参数均有 finite gradient。该结果验证 CP wrapper
和通信逻辑不依赖 Triton 才能正确工作。

5.3 Triton backend

配置	mode	output rel-L2	hidden grad rel-L2	最大参数 grad rel-L2	结果
CP4, global T=256, H=4	P2P	8.069e-3	8.878e-3	9.294e-3	PASS
CP4, global T=256, H=4	Ulysses	8.071e-3	8.736e-3	9.058e-3	PASS
CP8, global T=512, H=8	P2P	8.08e-3	8.87e-3	9.26e-3	PASS
CP8, global T=512, H=8	Ulysses	8.07e-3	8.84e-3	9.97e-3	PASS

CP4 与 CP8 的 P2P/Ulysses 均低于 3e-2 门槛，没有随 CP size 增大出现误差放大。

P2P 还单独覆盖随机非零 initial state 和随机非零 final-state gradient，并检查：

• local token output；
• rank final state；
• initial-state gradient；
• q/k/v/gate/beta/A_log/dt_bias 的全部梯度；
• forward affine summary 与 backward gradient-summary 的多 rank 顺序。

所有检查 PASS，未发现漏传 state gradient、summary 顺序反转或某个 rank 梯度缺失。

6. 多步 AdamW 训练轨迹

配置为 CP4,B=1,global T=256,hidden=256,H=4,Dk=Dv=128,BF16，AdamW 使用
lr=1e-3,weight_decay=0.01。Reference 与 CP 从相同权重和 optimizer state 开始；每一步
使用一组新的确定性随机输入和 target，共执行 20 个 optimizer step。

mode/backend	steps	最大 loss rel	最大 grad-norm rel	最终最大参数 rel-L2	结果
Ulysses/Eager	20	7.43e-6	4.79e-4	8.56e-3	PASS
P2P/Eager	20	8.67e-6	5.88e-4	8.54e-3	PASS
Ulysses/Triton	20	1.11e-5	6.00e-4	8.50e-3	PASS
P2P/Triton	20	8.74e-6	4.89e-4	8.36e-3	PASS

四种 mode/backend 组合均满足轨迹门槛。20 步内未出现 NaN/Inf、漏梯度、loss 突跳或
loss/grad-norm 误差随 step 持续扩大。Eager 与 Triton 的最终参数误差处于同一量级，未观察
到 P2P state-summary 反向带来的独有漂移。

7. 性能与显存

7.1 测试范围与配置

本章测试 KDA core 与对应 CP 通信，不包含 projection、ShortConv、output gate、RMSNorm 和
output projection。该口径用于直接比较 P2P state chain 与 Ulysses all-to-all，不代表完整
Kimi K3 模型的端到端加速比。

统一 shape：

B = 1
heads = 96
Dk = Dv = 128
chunk_size = 64
local sequence = 8K
dtype = BF16

7.2 CP4/CP8 结果

CP/global sequence	mode	Forward	Fwd+Bwd	Total peak allocated	Total peak reserved
CP4/32K	P2P	81.05 ms	228.38 ms	7748 MiB	9044 MiB
CP4/32K	Ulysses	90.26 ms	253.77 ms	7906 MiB	9250 MiB
CP8/64K	P2P	82.72 ms	231.81 ms	7748 MiB	9044 MiB
CP8/64K	Ulysses	89.04 ms	245.28 ms	7906 MiB	9254 MiB

对应收益：

配置	P2P Forward 加速	P2P Fwd+Bwd 加速	allocated 降低	reserved 降低
CP4/32K	10.2%	10.0%	2.0%	2.2%
CP8/64K	7.1%	5.5%	2.0%	2.3%

CP4 的 forward 阶段峰值如下：

mode	Forward peak allocated	Forward peak reserved
P2P	3119 MiB	3634 MiB
Ulysses	3854 MiB	4422 MiB

P2P 的 forward peak allocated/reserved 分别降低约 19.1%/17.8%。完整前反向峰值的差距
缩小，是因为两条路径都会进入占用更高的公共 KDA backward workspace；不能用 total peak
掩盖 forward 存活张量的差异。

保持 local sequence=8K 从 CP4 扩展到 CP8 后，P2P Forward 增加 2.1%，Fwd+Bwd 增加
1.5%，峰值 allocated 保持不变。结果说明该 shape 下 P2P 的 local 计算与显存主要由 local
sequence 决定。新增 rank 带来的状态链开销，以及本机 CP8 Ulysses 时间相对 CP4 略低，
共同使 CP8 的 P2P 相对加速比低于 CP4。

8. Activation Checkpoint

使用 non-reentrant activation checkpoint 验证 CP4 P2P 完整层。配置为
B=1,global T=256,H=4,BF16，开启和关闭 checkpoint 时使用相同权重、输入和上游梯度。

P2P 配置	output rel-L2	hidden grad rel-L2	最大参数 grad rel-L2	结果
checkpoint off	8.069e-3	8.878e-3	9.294e-3	PASS
checkpoint on	8.069e-3	8.878e-3	9.220e-3	PASS

Checkpoint 重放过程中未出现：

• forward P2P 消息序号复用；
• receive/send 次序错配；
• backward state-gradient 丢失；
• 死锁、NaN/Inf 或参数漏梯度。

开启与关闭 checkpoint 的误差处于相同量级，说明自定义 P2P autograd function 能在重计算
路径中重新建立独立通信上下文。

9. PR 可引用结论

1. Hyper-Parallel Eager KDA 已与 FLA recurrent/chunkwise reference 对齐。FP32 output、final
   state 和全部梯度最大绝对误差不超过 3.772e-7。
2. CP4 Eager affine summary 在随机非零 initial state 和 final-state gradient 下复现完整序列
   output、final state 和全部梯度，验证了 state-summary 正反向数学。
3. 单卡 Triton KDA 在 B=1,T=128,H=4,Dk=Dv=128,BF16 下通过全部前反向精度检查；
   最差 relative-L2 为 1.27e-2。
4. CP4/CP8 Triton Ulysses 与 P2P 完整 KimiDeltaAttention 层的 output、hidden gradient 和
   最大参数 gradient relative-L2 均低于 1.0e-2。
5. Eager 与 Triton、Ulysses 与 P2P 四种组合均完成 20 步非零学习率 AdamW 训练；最大 loss
   relative error 为 1.11e-5，最大 grad-norm relative error 为 6.00e-4，无持续漂移。
6. 正式 H=96,D=128,local sequence=8K KDA core 中，P2P 相对 Ulysses 的 Fwd+Bwd
   加速为 CP4 10.0%、CP8 5.5%。
7. P2P total peak allocated 比 Ulysses 低约 2.0%；CP4 forward peak allocated 低约
   19.1%。
8. CP4 P2P activation checkpoint 精度通过，重放期间未发生消息错配、死锁或梯度缺失。
9. Eager backend 不依赖 FLA/Triton；Triton backend 对依赖、dtype、shape 和 staged API
   执行显式 capability 检查，不支持条件不会静默回退。
10. KDA CPU UT、FLA adapter UT、pylint、compileall、Lizard 审计、wheel 构建和隔离安装均
    已通过。

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

The issue contains a Chinese test report for Hyper-Parallel PR #1311 and commit a867e401a287b6027d535c353c355ec85d72cd12, covering KDA, Ulysses, P2P, Eager, Triton, and checkpoint tests. Start by reviewing that PR and commit, then clarify the target documentation file and the specific change expected before implementing anything.