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 构建和隔离安装均
    已通过。