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

mhc_post 算子正确性证明与 AscendC 实现深度解析mHC 后连接层的广播缩放机制【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformermhc_post 是 CANN ops-transformer 仓库experimental/mhc/mhc_post目录下的一个 AscendC 算子它在 mHCMulti-head Hyper-Connections多头超连接框架中实现论文公式x_{l1} 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_{l1} 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 个 streamF(...)分支模块branch module输出branch_outputH_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]各张量的含义张量Shape说明branch_output[batch, seq_len, dim]分支模块 F(...) 的输出h_post[num_streams]可学习权重归一化在上游完成本算子不处理output[batch × num_streams, seq_len, dim]分发到 N 个 stream 的结果要点h_post是长度为num_streams的静态权重向量对所有 batch 与所有 token 位置共享权重归一化normalization由上游负责mhc_post 只做逐元素乘法与广播。2. PyTorch 参考实现tokenbender/mHCmhc_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×Ns, ...] for (int64_t i 0; i E; i) { // output[b×Ns, 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 Aper-stream逐流并行任务粒度 (batch, stream) 对每个任务完整读取输入行、乘以对应权重、写出一行输出。适合seq×dim较小/中等的场景。任务总数total_tasks batch × num_streams。Strategy Bread-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根据此规则选择内核并钳制 blockDimStrategy A 上限b×nStrategy 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×Ns, ...] 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队列完成DataCopyGM→UB→ Muls标量乘→ DataCopyUB→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.cppauto h_fp32 h_post.to(torch::kFloat32).contiguous()即h_post_fp32是 fp32 输入。该路径通过Cast→Muls→Cast三步完成牺牲少量额外 UB 空间临时 fp32 buffertmpBuf换取精度。5. 正确性映射论文 → PyTorch → CPU → NPU正确性文档用一张四层映射表把同一个数学运算在不同实现中的对应关系钉死Paper FormulaPyTorchCPUNPUbranch_output[b, ...]einsum input[B,S,D]branch_output[b * E i]gm_branch batch_idx * Eoutput[b×Ns, ...]rearrange(b s)output[(b*Ns) * E i]gm_output (b*Ns) * Eh_post[s]beta[s]h_post[s]gm_h.GetValue(stream_idx)× h_post[s]einsumb...d,s-b...sd* weightMuls(out, in, weight)这张表的价值在于四个实现虽然写法各异einsum 下标、双重循环、GlobalTensor 偏移、向量指令但共享完全相同的索引代数因此正确性可以逐项对应验证。6. 索引计算证明Index Calculation Proof正确性文档给出了严格的线性索引推导。给定Bbatch, Nnum_streams, Sseq_len, Ddim, ES×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×Ns, ...]由于in_off与out_off的差恒为s × E同 batch 下第 s 个 stream 的输出整体位于输入之后偏移s×E处且内层seq×D d部分完全一致因此逐元素乘法的位置映射是精确的——这是整个正确性论证的核心。7. 与 mhc_pre 的数学对偶关系mhc_post 与仓库中的 mhc_pre 算子在 mHC 框架中是数学对偶mathematical inverses关系Aspectmhc_postmhc_prePaper PartH_l^{post}^T · F(...)H_l^{pre} · x_lOperationBroadcast (1 → N)Reduce (N → 1)Input[B, S, D][B×N, S, D]Output[B×N, S, D][B, S, D]Formulaout[b×Ns] in[b] × w[s]out[b] Σ_s in[b×Ns] × w[s]PyTorchb...d, s - b...sdbs...d, s - b...dblockDimB × NB直观理解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_exactyes PASS (0 mismatch) FP16: max_abs1.22e-04 PASS BF16: max_abs9.73e-04 PASS三种精度采用不同判定标准dtypePrecision Criterionfp32bit-exact逐位比较ULP0bit_copy后直接比对 uint32fp16allclose(atol1e-4, rtol1e-3)bf16allclose(atol1e-3, rtol4e-3)bf16 尾数仅 8 位容差放宽FP32 的 bit-exact 成立是有内在原因的Muls(out, in, weight)与 CPU 上的branch_output[i] * weight遵循相同的 IEEE-754 浮点乘语义单次乘法的舍入结果应完全一致因此可以实现 0 mismatch 的逐位一致。8.2 边界用例test_edge_cases.cpp 针对非对齐维度与极端形状 Edge Cases dim1, dim7, dim15 PASS (non-aligned) batch1, seq1 PASS (boundary) num_streams1,2,4,8 PASS (various N)dim1/7/15验证非 8/16 对齐的维度fp32 对齐 ALIGN8fp16/bf16 对齐 ALIGN16在DataCopyPad尾部填充路径下结果正确batch1, seq1验证最小边界形状num_streams 1/2/4/8验证不同 stream 数上限 8 与内核中MAX_STREAMS 8一致。9. 构建、运行与使用9.1 构建参照 README.mdsource /usr/local/Ascend/ascend-toolkit/set_env.sh # 1. Build AscendC kernel mkdir -p build cd build cmake .. -DSOC_VERSIONascend910b2 make -j cd .. # 2. Build PyTorch C extension python setup.py build_ext --inplace9.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.py9.3 Python APIPyTorch 侧通过 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, dtypetorch.float32, devicenpu) h torch.randn(N, dtypetorch.float32, devicenpu) out mhc_post_ext.forward(x, h) # [B*N, S, D]或使用封装层 mhc_post_ops.pyfrom 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 910B220 AI Core上相对torch.einsumNPU 上执行的对比多数 shape 获得 2~4 倍加速如(4,512,256) ns4为 3.8x、(8,256,512) ns4为 3.9x个别batch16大元素 shape 略慢文档说明是框架层 aclnnMul 优化所致。内核实现要点包括 192KB UB_SIZE、双缓冲BUFFER_NUM2、按seq×dim与 dtype 动态计算 tile 大小。10. 总结mhc_post 是 mHC 深度连接公式中后连接post-connection的 NPU 落地实现其正确性由三层证据链保证语义层与论文公式output[b×Ns, seq, d] branch_output[b, seq, d] × h_post[s]及 PyTorch einsum/rearrange 参考实现逐项对应索引层线性索引推导证明 NPU 内核的in_off/out_off偏移与 CPU 参考完全一致验证层fp32 逐位一致bit-exact、fp16/bf16 按精度容差通过配合非对齐维度与极端形状边界用例。配合 mhc_pre 构成数学对偶的归约-广播配对为 mHC 这类多 stream 深度连接架构在昇腾 NPU 上提供了高性能、可验证的基础算子。感兴趣的读者可继续阅读 mhc_pre 文档 对照学习其对偶实现。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考