企业官网建设流程全解析

mhc_post 算子代码级正确性证明：从 mHC 论文公式到 PyTorch 参考实现与 AscendC NPU Kernel 的全链路推导

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

mhc_post 是 CANN ops-transformer 仓库中 mHC（Manifold-Constrained Hyper-Connections）注意力模型在昇腾 NPU 上的关键后处理算子，负责将分支模块（Attention / FFN）的输出按H_post权重分发到多个流（streams）。本文以 v1 原型仓库 中的代码级正确性证明文档为核心，逐层还原「论文公式 → PyTorch 参考实现 → CPU 参考实现 → NPU Kernel 实现 → 数学等价性证明」的完整推导链路，并结合 Kernel 源码与测试用例，说明每一层实现如何相互印证，最终证明 NPU 实现与论文公式完全等价。

1. 背景：mhc_post 在 mHC 架构中的位置

mHC 论文（arXiv:2512.24880，DeepSeek 2024.12.31 发布）提出了流形约束超连接（Manifold-Constrained Hyper-Connections）架构。其核心迭代公式（Equation 3）为：

x_{l+1} = H_l^{res} · x_l + H_l^{post}^T · F(H_l^{pre} · x_l, W_l)

各符号含义如下：

从仓库中的 验证脚本 可以看到，mhc_post 算子对应的正是公式中的H_l^{post}^T · F(...)部分：将分支模块的输出F(...)与经过 softmax 归一化的H_post向量做外积广播，再分发到多个 streams。

2. 论文公式的精确定义

mhc_post 的论文公式可写成如下一行：

mhc_post: output[b*s+i, seq, d] = branch_output[b, seq, d] × h_post[i]

其中输入输出形状为：

• branch_output: [batch, seq_len, dim] —— 分支模块（如 Attention/FFN）的输出
• h_post: [num_streams] —— 分发权重，即H_l^{post}，经 softmax 归一化
• output: [batch * num_streams, seq_len, dim] —— 每个 (batch, stream) 组合一份输出

公式要义是：对每一个 batch 样本，将其同一份分支输出复制到 num_streams 份，每份乘以对应 stream 的权重h_post[s]，最终按batch * num_streams + stream的次序把 batch 与 stream 两个维度合并。这个算子不涉及任何累加或规约，只有「广播乘法 + 维度重排」两个动作，因此正确性证明可以做到逐元素级别。

3. PyTorch 参考实现（tokenbender/mHC）

仓库文档给出的 PyTorch 参考实现来自hyper_connections_mhc.py的depth_connection()函数，它用两步张量操作完成上述公式：

def depth_connection(self, branch_output, residuals, *, beta): # beta 就是 h_post，shape [num_streams] # Step 1: einsum 广播乘法 # "b ... d, s -> b ... s d" 的含义: # - 输入1 branch_output: [batch, seq, dim] # - 输入2 beta: [streams] # - 输出: [batch, seq, streams, dim] # - 计算: out[b,seq,s,d] = branch_output[b,seq,d] * beta[s] output = einsum(branch_output, beta, "b ... d, s -> b ... s d") # Step 2: reshape 合并 batch 和 stream 维度 # "b ... s d -> (b s) ... d" 的含义: # - 输入: [batch, seq, streams, dim] # - 输出: [batch*streams, seq, dim] # - 计算: out[b*s+i, seq, d] = in[b, seq, i, d] output = rearrange(output, "b ... s d -> (b s) ... d") return output

这两步与论文公式的对应关系非常清晰：

• einsum"b ... d, s -> b ... s d"在 batch 与 stream 之间没有求和（einsum 输出中同时保留 b 与 s 下标），本质是外积广播：output[b,seq,s,d] = branch_output[b,seq,d] × beta[s]；
• rearrange"b ... s d -> (b s) ... d"把[batch, seq, streams, dim]按「先 batch 后 stream」的次序压平为[batch*streams, seq, dim]，即output[b*streams + s, seq, d] = input[b, seq, s, d]。

两步组合的结果正是论文公式。仓库中的 Python 验证脚本 用np.einsum('bsd,n->bsnd', ...)加 transpose/reshape 复现了该参考实现，并断言与参考实现逐元素最大差异小于 1e-6。

4. CPU 参考实现（测试代码中的mhc_post_cpu）

在编写 NPU Kernel 之前，原型先提供了一个可直接对照的 CPU 参考实现mhc_post_cpu()，完整代码位于 test_mhc_post.cpp：

void mhc_post_cpu( const float* branch_output, // [batch, seq_len, dim] const float* h_post, // [num_streams] float* output, // [batch * num_streams, seq_len, dim] int64_t batch, int64_t seq_len, int64_t dim, int64_t num_streams ) { // 遍历每个batch for (int64_t b = 0; b < batch; ++b) { // 遍历每个stream for (int64_t s = 0; s < num_streams; ++s) { // 获取当前stream的权重 float weight = h_post[s]; // 计算输出的batch索引: b*num_streams + s // 这对应 rearrange "b s -> (b s)" 的语义 int64_t out_batch_idx = b * num_streams + s; // 遍历每个序列位置和维度 for (int64_t seq = 0; seq < seq_len; ++seq) { for (int64_t d = 0; d < dim; ++d) { // 输入索引: branch_output[b, seq, d] int64_t in_idx = b * seq_len * dim + seq * dim + d; // 输出索引: output[b*s+i, seq, d] int64_t out_idx = out_batch_idx * seq_len * dim + seq * dim + d; // 核心计算: 乘法 // 对应论文公式: output[b*s+i] = branch_output[b] * h_post[i] output[out_idx] = branch_output[in_idx] * weight; } } } } }

这份实现的价值在于：它把论文公式中张量层面的「广播乘法 + 维度合并」显式翻译成了标量循环，每一行代码都可以与公式逐项对应：

CPU 实现完全按照行主序（row-major）线性索引展开，seq * dim + d对应[seq, d]平面的连续排布，b * seq_len * dim对应 batch 维的步长，out_batch_idx * seq_len * dim对应合并后 batch 维的步长，与 PyTorch 的[batch*streams, seq, dim]布局一致。

5. NPU Kernel 实现（AscendCMhcPostKernel）

NPU 端使用 AscendC 编程模型实现，完整代码位于 kernel/mhc_post_kernel.cpp。核心思路是把每个 (batch, stream) 组合映射为一个 block，从而天然并行地覆盖batch * num_streams个输出切片。

5.1 Init：block 到 (batch, stream) 的索引解析

__aicore__ inline void Init( GM_ADDR branch_output, GM_ADDR h_post, GM_ADDR output, int64_t batch, int64_t seq_len, int64_t dim, int64_t num_streams ) { this->batch = batch; this->seq_len = seq_len; this->dim = dim; this->num_streams = num_streams; int64_t block_idx = GetBlockIdx(); this->batch_idx = block_idx / num_streams; // 解析 batch 索引 this->stream_idx = block_idx % num_streams; // 解析 stream 索引 this->batch_elements = seq_len * dim; // 每个 batch 的元素数 // 输入偏移: 只依赖 batch_idx，对应 branch_output[b, ...] int64_t input_offset = this->batch_idx * this->batch_elements; this->gm_branch_output.SetGlobalBuffer( reinterpret_cast<__gm__ float*>(branch_output) + input_offset, this->batch_elements ); this->gm_h_post.SetGlobalBuffer( reinterpret_cast<__gm__ float*>(h_post), num_streams ); // 输出偏移: 依赖 batch_idx 和 stream_idx，对应 output[b*s+i, ...] int64_t output_offset = (this->batch_idx * num_streams + this->stream_idx) * this->batch_elements; this->gm_output.SetGlobalBuffer( reinterpret_cast<__gm__ float*>(output) + output_offset, this->batch_elements ); this->tile_length = 256; // 固定 tile 大小 this->tile_num = (this->batch_elements + this->tile_length - 1) / this->tile_length; pipe.InitBuffer(inQueue, BUFFER_NUM, this->tile_length * sizeof(float)); pipe.InitBuffer(outQueue, BUFFER_NUM, this->tile_length * sizeof(float)); }

其中GetBlockIdx()返回当前 block 编号。由于启动的 block 总数等于batch * num_streams，block_idx / num_streams与block_idx % num_streams恰好完成论文公式中(b s) -> b, s的反向解析，对应output[b*s+i]中索引的拆解。

5.2 LoadWeight：加载当前 stream 权重

__aicore__ inline void LoadWeight() { // 加载当前stream的权重，对应论文公式中 h_post[i] this->weight_value = this->gm_h_post.GetValue(this->stream_idx); }

h_post只有num_streams个标量，每个 block 只需按自己的stream_idx读取一个权重即可。

5.3 Compute：向量乘标量

__aicore__ inline void Compute(int64_t length) { LocalTensor<float> inLocal = inQueue.DeQue<float>(); LocalTensor<float> outLocal = outQueue.AllocTensor<float>(); // 核心计算: 向量乘标量，对应论文公式: output = branch_output * h_post[i] // Muls 是 AscendC 的向量乘标量指令 Muls(outLocal, inLocal, this->weight_value, length); outQueue.EnQue(outLocal); inQueue.FreeTensor(inLocal); }

NPU 端用一条Muls（向量乘标量）指令完成整个[seq_len, dim]平面（或其中一段 tile）的逐元素乘法，与 CPU 参考实现中的标量乘法循环语义一致。由于 mhc_post 只需要乘法、不涉及累加/除法等复杂运算，FP32 下理论上可以达到与 CPU 参考实现 bit-exact 的精度（见第 7 节验证）。

5.4 Process：tiling 流水

__aicore__ inline void Process() { this->weight_value = this->gm_h_post.GetValue(this->stream_idx); for (int64_t tile_idx = 0; tile_idx < this->tile_num; ++tile_idx) { int64_t offset = tile_idx * this->tile_length; int64_t length = this->tile_length; if (offset + length > this->batch_elements) { length = this->batch_elements - offset; // 尾块裁剪 } CopyIn(offset, length); Compute(length); CopyOut(offset, length); } }

batch_elements = seq_len * dim可能很大，Kernel 用固定的tile_length = 256对每个 (batch, stream) 的[seq_len, dim]平面做分块，配合BUFFER_NUM = 2的双缓冲队列（inQueue/outQueue）实现「搬入-计算-搬出」的流水。尾块（不整除 256）通过length = batch_elements - offset精确裁剪。

5.5 CopyIn / CopyOut：对齐与非对齐路径

constexpr int32_t ALIGN_ELEM = 8; // fp32 对齐元素数 __aicore__ inline void CopyIn(int64_t offset, int64_t length) { LocalTensor<float> inLocal = inQueue.AllocTensor<float>(); if (length % ALIGN_ELEM == 0) { DataCopy(inLocal, gm_branch_output[offset], length); } else { DataCopyExtParams params{1, static_cast<uint32_t>(length * sizeof(float)), 0, 0, 0}; DataCopyPadExtParams<float> pad{false, 0, 0, 0.0f}; DataCopyPad(inLocal, gm_branch_output[offset], params, pad); } inQueue.EnQue(inLocal); } __aicore__ inline void CopyOut(int64_t offset, int64_t length) { LocalTensor<float> outLocal = outQueue.DeQue<float>(); if (length % ALIGN_ELEM == 0) { DataCopy(gm_output[offset], outLocal, length); } else { DataCopyExtParams params{1, static_cast<uint32_t>(length * sizeof(float)), 0, 0, 0}; DataCopyPad(gm_output[offset], outLocal, params); } outQueue.FreeTensor(outLocal); }

当length是 8 的倍数（fp32 下的 32 字节对齐）时走常规DataCopy；否则走DataCopyPad/DataCopyExtParams的带填充路径，保证任意seq_len * dim（包括非对齐的 dim，如dim=7）都能正确处理。这与仓库中 test_unaligned.cpp 针对dim=7非对齐场景的测试一一对应。

5.6 Kernel 入口与 do 接口

extern "C" __global__ __aicore__ void mhc_post_kernel( GM_ADDR branch_output, GM_ADDR h_post, GM_ADDR output, int64_t batch, int64_t seq_len, int64_t dim, int64_t num_streams ) { MhcPostKernel op; op.Init(branch_output, h_post, output, batch, seq_len, dim, num_streams); op.Process(); } extern "C" void mhc_post_do( uint32_t blockDim, void* stream, uint8_t* branch_output, uint8_t* h_post, uint8_t* output, int64_t batch, int64_t seq_len, int64_t dim, int64_t num_streams ) { mhc_post_kernel<<<blockDim, nullptr, stream>>>( branch_output, h_post, output, batch, seq_len, dim, num_streams ); }

mhc_post_kernel是 AscendC 的__global__ __aicore__入口，mhc_post_do是对外的可调用封装：调用方传入blockDim（即batch * num_streams）与 ACL stream，通过<<<blockDim, nullptr, stream>>>启动核函数。

5.7 NPU 实现与论文公式的逐项对照

6. 数学等价性证明

论文公式、CPU 实现与 NPU 实现三者之间的等价性可以形式化如下。

给定：

• branch_output: shape [B, S, D]
• h_post: shape [N]
• output: shape [B*N, S, D]

论文公式：

output[b*N + n, s, d] = branch_output[b, s, d] × h_post[n] 其中 b ∈ [0,B), n ∈ [0,N), s ∈ [0,S), d ∈ [0,D)

CPU 实现等价性：

for b in [0, B): for n in [0, N): out_idx = b * N + n // ← 对应 output[b*N + n, ...] weight = h_post[n] // ← 对应 h_post[n] for s in [0, S): for d in [0, D): output[out_idx, s, d] = branch_output[b, s, d] * weight // ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ // 对应 branch_output[b, s, d] × h_post[n]

NPU 实现等价性：

// 并行化: 每个 block 处理一个 (b, n) 组合 block_idx ∈ [0, B*N): b = block_idx / N // ← 解析 batch 索引 n = block_idx % N // ← 解析 stream 索引 weight = h_post[n] // ← 对应 h_post[n] // 向量化: 处理整个 [S, D] 平面 output[block_idx, :, :] = branch_output[b, :, :] * weight // ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ // 对应 branch_output[b, s, d] × h_post[n]

由于block_idx = b * N + n，所以output[block_idx]等价于output[b*N + n]；因此NPU 实现与论文公式完全等价。

7. 测试用例与精度验证

正确性证明文档之外，原型目录下的测试代码从多个维度对该等价性做了工程验证：

7.1 基础正确性测试

test_mhc_post.cpp 是核心对照测试：以batch=2, seq_len=4, dim=64, num_streams=4为例，h_post先按[1,2,3,4]初始化再经softmax归一化，branch_output用(i % 100) / 100.0f填充；先跑 CPU 参考实现得到期望输出，再通过 ACL 接口分配设备内存、以blockDim = batch * num_streams启动mhc_post_do，最后以1e-4容差逐元素比对 CPU 与 NPU 结果。该文件同时给出了softmax的 C++ 实现（先减最大值再指数归一化），与h_post需要满足 softmax 约束的论文设定一致。

7.2 多 shape 综合测试

test_multi_shape.cpp 覆盖 10 组形状：基础测试、单 batch、大 batch、长序列、大 dim、多 streams、典型配置、dim=32边界、streams=2边界以及大规模（4×64×256×4），全部使用随机正态分布初始化并经 softmax，判定条件为最大绝对误差< 1e-4且最大相对误差< 1e-3。

7.3 非对齐场景测试

test_unaligned.cpp 专门针对dim=7（元素数与 8 不对齐）验证 Kernel 的DataCopyPad分支，容差1e-5。

7.4 精度详细测试

test_precision.cpp 用 5 组不同输入尺度（小数值 0.01、标准正态、大数值 100、混合正负、大规模）统计最大绝对/相对误差与完全一致点数，判定阈值max_abs < 1e-5。

7.5 FP32 精度极限分析

test_fp32_limit.cpp 对比了 FP32 理论极限（machine epsilon ≈ 1.19e-7，约 7 位十进制有效数字）与 6 个数量级范围（从 [1e-30, 1e-20] 到 [1e20, 1e30] 及混合正负）的实际测试结果，结论是mhc_post 的 Muls 乘法在所有测试范围内与 CPU 参考实现达到 bit-exact——因为昇腾 NPU 的 FP32 乘法遵循 IEEE 754 标准，而单次乘法不引入累加误差。

7.6 Python 侧公式验证

verify_mhc_post.py 从公式层面对三类结论做交叉验证：

• verify_formula()：小规模（batch=2, seq_len=2, dim=4, num_streams=3）手工逐点验证output[out_idx,0,0] = branch_output[b,0,0] * h_post[s]；
• compare_with_einsum()：将参考实现与np.einsum('bsd,n->bsnd')+ transpose/reshape 对比，断言最大差异< 1e-6，等价于 PyTorch 的 einsum + rearrange 参考路径；
• generate_test_vectors()：按与 C++ 测试完全一致的初始化方式（(i % 100)/100、h_post=[1,2,3,4]softmax 归一化）生成二进制测试向量，供 NPU 端对照。

其余两个脚本 test_comprehensive.py 与 test_formula_proof.py 进一步补充了综合场景与公式级证明。test/下还有 test_mhc_post_detail.cpp 提供更细粒度的逐点输出检查。

8. 构建与运行环境说明

原型目录提供 CMakeLists.txt 作为构建入口，其关键配置如下：

• CMAKE_CXX_STANDARD 17，项目名mhc_post_v1；
• 通过include(${ASCEND_PATH}/compiler/tikcpp/ascendc_kernel_cmake/ascendc.cmake)引入 AscendC 编译框架（ASCEND_PATH默认指向/usr/local/Ascend/ascend-toolkit/latest）；
• SOC_VERSION默认为ascend910b2，可通过 CMake 缓存变量覆盖；
• ascendc_library(mhc_post_kernel STATIC kernel/mhc_post_kernel.cpp)将 Kernel 编译为静态库；
• 测试程序通过宏add_test_exe生成，链接mhc_post_kernel ascendcl runtime并引入${ASCEND_PATH}/include。

测试程序运行依赖 ACL（AscendCL）运行时：需要先aclInit、aclrtSetDevice(0)、aclrtCreateStream，通过aclrtMalloc/aclrtMemcpy完成主机与设备间的数据搬运，最后aclrtSynchronizeStream同步后再把设备端结果拷回主机比对。v1 原型为仅 fp32、固定 tile=256 的初始原型，如 README 所注，已归档仅供后续版本参考。

9. 小结：四层实现如何相互印证

mhc_post 的代码级正确性证明，本质是一条「从数学到机器指令」的等值链：

1. 论文公式定义语义：output[b*N + n, s, d] = branch_output[b, s, d] × h_post[n]；
2. PyTorch 参考实现（einsum + rearrange）把公式翻译为张量操作，两步分别对应「外积广播」与「维度合并」；
3. CPU 参考实现（mhc_post_cpu）把张量操作再翻译为逐元素标量循环，成为可逐项对照的黄金基准；
4. NPU Kernel 实现（MhcPostKernel）通过block_idx / num_streams、block_idx % num_streams复现索引解析，用Muls完成乘法，用 tiling + 双缓冲 + 对齐/填充两条 DMA 路径覆盖任意形状，最终由block_idx = b * N + n的构造保证与论文公式逐元素等价。

这一「公式 → 框架参考实现 → CPU 基准 → 硬件 Kernel → 数学证明 + 多维测试」的方法论，也为后续在昇腾 NPU 上迁移其他 mHC 算子（如mhc_pre、mhc_sinkhorn等，见 mhc 目录）提供了可复制的验证范式。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer