企业官网建设流程全解析

mhc_post 算子正确性证明与 AscendC 实现深度解析：mHC 后连接层的广播缩放机制

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

mhc_post 是 CANN ops-transformer 仓库experimental/mhc/mhc_post目录下的一个 AscendC 算子，它在 mHC（Multi-head Hyper-Connections，多头超连接）框架中实现论文公式x_{l+1} = H_l^{res}·x_l + H_l^{post}^T·F(H_l^{pre}·x_l, W_l)中的H_l^{post}^T·F(...)部分——即将分支模块输出的 1 个 stream 按可学习权重广播扩展到 N 个 stream。本文以 proof_of_correctness.md 为骨架，结合 内核实现、测试用例 与 PyTorch 封装，完整梳理其论文公式、参考实现、NPU 内核、索引映射证明与验证体系，帮助读者理解该算子"为什么正确"以及"如何在 NPU 上高效落地"。

1. 算子在 mHC 框架中的定位

1.1 mHC 论文公式

mHC 论文（Equation 3）给出了深度连接的递推形式：

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

其中：

• H_l^{pre} · x_l：由 mhc_pre 算子负责，将 N 个 stream 归约（Reduce）为 1 个 stream；
• F(...)：分支模块（branch module），输出branch_output；
• H_l^{post}^T · F(...)：由mhc_post负责，将 1 个输入广播（Broadcast）为 N 个 stream，并逐 stream 缩放。

mhc_post 计算的正是公式中标注的部分：将 1 个输入广播到 N 个 stream。

1.2 核心计算公式

output[b × N + s, seq, d] = branch_output[b, seq, d] × h_post[s]

各张量的含义：

要点：h_post是长度为num_streams的静态权重向量，对所有 batch 与所有 token 位置共享；权重归一化（normalization）由上游负责，mhc_post 只做逐元素乘法与广播。

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

mhc_post 的数学语义与开源实现 tokenbender/mHC 的depth_connection()完全一致（仓库 README.md 明确说明这一点）。其核心逻辑如下：

# Source: hyper_connections_mhc.py depth_connection() def depth_connection(self, branch_output, residuals, *, beta): # beta is h_post, shape [num_streams] # branch_output shape: [batch, seq, dim] # Step 1: "b ... d, s -> b ... s d" # Broadcast multiply: [B, S, D] × [N] -> [B, S, N, D] # 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: "b ... s d -> (b s) ... d" # Reshape: [B, S, N, D] -> [B×N, S, D] output = rearrange(output, "b ... s d -> (b s) ... d") return output

两步语义：先用 einsum 完成"广播乘"（branch_output[b,seq,d] × beta[s]），再通过 rearrange 把新增的 stream 维合并进 batch 维，得到[B×N, S, D]的输出。

仓库中的纯 Python 参考实现 mhc_post_ops.py 用一行等价 einsum 表达同一语义：

def mhc_post_einsum(x: torch.Tensor, h_post: torch.Tensor) -> torch.Tensor: batch = x.size(0) num_streams = h_post.size(0) if num_streams == 0: raise ValueError("num_streams must be > 0, got 0") return torch.einsum('bsd,n->bnsd', x, h_post).reshape(batch * num_streams, -1, x.size(-1))

3. CPU 参考实现：线性索引视角

为了在 NPU 上验证正确性，仓库维护了 CPU 参考实现（见 test_multi_dtype.cpp 中的cpu_reference_fp32，正确性文档中同样给出）：

void cpu_reference_fp32( const float* branch_output, // [B, S, D] const float* h_post, // [N] float* output, // [B×N, S, D] int64_t batch, int64_t seq_len, int64_t dim, int64_t num_streams ) { int64_t E = seq_len * dim; // elements per batch for (int64_t b = 0; b < batch; ++b) { for (int64_t s = 0; s < num_streams; ++s) { float weight = h_post[s]; int64_t out_batch = b * num_streams + s; // output[b×N+s, ...] for (int64_t i = 0; i < E; ++i) { // output[b×N+s, i] = branch_output[b, i] × h_post[s] output[out_batch * E + i] = branch_output[b * E + i] * weight; } } } }

该实现把 batch 与 stream 两层循环外提，内层对E = seq_len × dim个元素做标量乘法，是验证 NPU 内核的"黄金标准"。测试中以固定随机种子构造输入、将权重归一化（h_weight[i] /= sum）后比较。

4. NPU AscendC 内核实现

4.1 双策略自适应调度

内核 mhc_post_kernel.cpp 采用自适应策略，根据 shape 在两个并行化方案间自动切换：

• Strategy A（per-stream，逐流并行）：任务粒度 = (batch, stream) 对，每个任务完整读取输入行、乘以对应权重、写出一行输出。适合seq×dim较小/中等的场景。任务总数total_tasks = batch × num_streams。
• Strategy B（read-once，一次读取）：任务粒度 = (batch, tile) 对，每个任务只读一次输入 tile，连续写出 N 份（每个 stream 一份）输出。适合seq×dim大且 batch 大的场景。任务总数total_tasks = batch × tiles_per_batch。

策略选择规则在内核中体现为 UseReadOnce 函数：

constexpr int64_t READONCE_THRESHOLD_BYTES = 4 * 1024 * 1024; inline bool UseReadOnce(int64_t batch_elements, int64_t elem_size, int64_t num_streams, int64_t batch) { int64_t total_read = batch_elements * elem_size * num_streams; return total_read >= READONCE_THRESHOLD_BYTES && batch >= 16; }

即：当「seq×dim × sizeof(T) × num_streams ≥ 4MB」且「batch ≥ 16」时启用 Strategy B，因为此时对输入重复读取 N 次的代价过高，一次读取、多次写入更具优势。Host 侧入口mhc_post_do_fp32/fp16/bf16根据此规则选择内核，并钳制 blockDim（Strategy A 上限b×n，Strategy B 上限取b×tiles与 20 的较小者，见 内核入口）。

4.2 Strategy A 内核伪代码

正确性文档给出了内核的核心处理逻辑：

__aicore__ inline void ProcessOne(int64_t batch_idx, int64_t stream_idx) { // in_off = b × E ← branch_output[b, ...] int64_t in_off = batch_idx * batch_elements; // out_off = (b × N + s) × E ← output[b×N+s, ...] int64_t out_off = (batch_idx * num_streams + stream_idx) * batch_elements; gm_in.SetGlobalBuffer(gm_branch + in_off, batch_elements); gm_out.SetGlobalBuffer(gm_output + out_off, batch_elements); T weight = gm_h.GetValue(stream_idx); // h_post[s] for (int64_t i = 0; i < tiles; ++i) { CopyIn(off, len); Compute(len, weight); // Muls(out, in, weight) CopyOut(off, len); } }

实际内核代码在此基础上加入了双缓冲（BUFFER_NUM = 2）流水线：ProcessTile内部使用inQue/outQue队列完成DataCopy（GM→UB）→ Muls（标量乘）→ DataCopy（UB→GM）的搬运-计算-搬出循环，并对非对齐长度（l % ALIGN != 0）使用DataCopyPad做尾部填充处理。

4.3 BF16 的特殊处理路径

BF16 场景下，由于 bf16 尾数只有 8 位，直接做乘法的精度不足，内核为 BF16 单独实现了MhcPostPerStreamBF16/MhcPostReadOnceBF16，采用fp32 计算路径：

Cast(tmp, in, RoundMode::CAST_NONE, aligned); // bf16 -> fp32 Muls(tmp, tmp, weight, aligned); // fp32 标量乘 Cast(out, tmp, RoundMode::CAST_RINT, aligned); // fp32 -> bf16（就近舍入）

注意 BF16 的权重在 Host 侧被显式升为 fp32（见 mhc_post_torch.cpp：auto h_fp32 = h_post.to(torch::kFloat32).contiguous()），即h_post_fp32是 fp32 输入。该路径通过Cast→Muls→Cast三步完成，牺牲少量额外 UB 空间（临时 fp32 buffertmpBuf）换取精度。

5. 正确性映射：论文 → PyTorch → CPU → NPU

正确性文档用一张四层映射表把同一个数学运算在不同实现中的对应关系钉死：

这张表的价值在于：四个实现虽然写法各异（einsum 下标、双重循环、GlobalTensor 偏移、向量指令），但共享完全相同的索引代数，因此正确性可以逐项对应验证。

6. 索引计算证明（Index Calculation Proof）

正确性文档给出了严格的线性索引推导。给定B=batch, N=num_streams, S=seq_len, D=dim, E=S×D：

论文要求output[b×N + s, seq, d] = branch_output[b, seq, d] × h_post[s]，按行主序展开线性索引：

branch_output[b, seq, d] → b × E + seq × D + d ✓ output[b×N + s, seq, d] → (b×N + s) × E + seq × D + d ✓

NPU 内核中的偏移计算与之一一对应：

in_off = batch_idx * batch_elements = b × E // ✓ matches branch_output[b, ...] out_off = (batch_idx * num_streams + stream_idx) * batch_elements = (b × N + s) × E // ✓ matches output[b×N+s, ...]

由于in_off与out_off的差恒为s × E（同 batch 下第 s 个 stream 的输出整体位于输入之后偏移s×E处），且内层seq×D + d部分完全一致，因此逐元素乘法的位置映射是精确的——这是整个正确性论证的核心。

7. 与 mhc_pre 的数学对偶关系

mhc_post 与仓库中的 mhc_pre 算子在 mHC 框架中是数学对偶（mathematical inverses）关系：

直观理解：mhc_pre 把 N 个 stream 加权求和压回 1 个，mhc_post 把 1 个 stream 加权复制成 N 个，二者组合即构成深度连接的前后两半。从并行度上看，mhc_post 的任务规模天然是 mhc_pre 的 N 倍（blockDim 取B×NvsB），这也是其需要精细化任务切分的原因。

8. 测试验证体系

8.1 多精度测试

test_multi_dtype.cpp 对 fp32/fp16/bf16 三种精度分别测试（形状覆盖(2,64,256)×4与(4,32,128)×8），正确性文档给出的通过标准与实测结果：

=== mhc_post Multi-DType Test === FP32: bit_exact=yes PASS (0 mismatch) FP16: max_abs=1.22e-04 PASS BF16: max_abs=9.73e-04 PASS

三种精度采用不同判定标准：

FP32 的 bit-exact 成立是有内在原因的：Muls(out, in, weight)与 CPU 上的branch_output[i] * weight遵循相同的 IEEE-754 浮点乘语义，单次乘法的舍入结果应完全一致，因此可以实现 0 mismatch 的逐位一致。

8.2 边界用例

test_edge_cases.cpp 针对非对齐维度与极端形状：

=== Edge Cases === dim=1, dim=7, dim=15 PASS (non-aligned) batch=1, seq=1 PASS (boundary) num_streams=1,2,4,8 PASS (various N)

• dim=1/7/15：验证非 8/16 对齐的维度（fp32 对齐 ALIGN=8，fp16/bf16 对齐 ALIGN=16）在DataCopyPad尾部填充路径下结果正确；
• batch=1, seq=1：验证最小边界形状；
• num_streams 1/2/4/8：验证不同 stream 数（上限 8 与内核中MAX_STREAMS = 8一致）。

9. 构建、运行与使用

9.1 构建

参照 README.md：

source /usr/local/Ascend/ascend-toolkit/set_env.sh # 1. Build AscendC kernel mkdir -p build && cd build cmake .. -DSOC_VERSION=ascend910b2 make -j cd .. # 2. Build PyTorch C++ extension python setup.py build_ext --inplace

9.2 测试

# C++ test cd build && LD_LIBRARY_PATH=./lib:$LD_LIBRARY_PATH ./test_multi_dtype # Python test LD_LIBRARY_PATH=./build/lib:$LD_LIBRARY_PATH python mhc_post_ops.py

9.3 Python API

PyTorch 侧通过 C++ 扩展 mhc_post_torch.cpp 暴露forward接口：mhc_post_forward校验输入连续性与num_streams > 0，按 dtype 分发到 fp32/fp16/bf16 内核（bf16 权重先转 fp32），并在当前 NPU stream 上以block_dim（默认 0，由内核自动选择）启动。使用方式：

import mhc_post_ext x = torch.randn(B, S, D, dtype=torch.float32, device='npu') h = torch.randn(N, dtype=torch.float32, device='npu') out = mhc_post_ext.forward(x, h) # [B*N, S, D]

或使用封装层 mhc_post_ops.py：

from mhc_post_ops import mhc_post, mhc_post_einsum out = mhc_post(x, h) # NPU 内核路径 ref = mhc_post_einsum(x, h) # einsum 参考路径，可交叉验证

// C++ kernel entry（自动选择策略） extern "C" void mhc_post_do_fp32(uint32_t blockDim, void* stream, uint8_t* input, uint8_t* h_post, uint8_t* output, int64_t batch, int64_t seq_len, int64_t dim, int64_t num_streams);

9.4 性能参考

performance.md 记录了在 Ascend 910B2（20 AI Core）上相对torch.einsum（NPU 上执行）的对比，多数 shape 获得 2~4 倍加速（如(4,512,256) ns=4为 3.8x、(8,256,512) ns=4为 3.9x）；个别batch=16大元素 shape 略慢，文档说明是框架层 aclnnMul 优化所致。内核实现要点包括 192KB UB_SIZE、双缓冲BUFFER_NUM=2、按seq×dim与 dtype 动态计算 tile 大小。

10. 总结

mhc_post 是 mHC 深度连接公式中"后连接（post-connection）"的 NPU 落地实现，其正确性由三层证据链保证：

1. 语义层：与论文公式output[b×N+s, seq, d] = branch_output[b, seq, d] × h_post[s]及 PyTorch einsum/rearrange 参考实现逐项对应；
2. 索引层：线性索引推导证明 NPU 内核的in_off/out_off偏移与 CPU 参考完全一致；
3. 验证层：fp32 逐位一致（bit-exact）、fp16/bf16 按精度容差通过，配合非对齐维度与极端形状边界用例。

配合 mhc_pre 构成数学对偶的"归约-广播"配对，为 mHC 这类多 stream 深度连接架构在昇腾 NPU 上提供了高性能、可验证的基础算子。感兴趣的读者可继续阅读 mhc_pre 文档 对照学习其对偶实现。

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