ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

CANN ops-nn:LayerNormStride 跨步 LayerNorm 算子——面向 MLA 场景的参数说明、调用方式与 Ascend C 内核实现剖析

CANN ops-nn:LayerNormStride 跨步 LayerNorm 算子——面向 MLA 场景的参数说明、调用方式与 Ascend C 内核实现剖析 CANN ops-nnLayerNormStride 跨步 LayerNorm 算子——面向 MLA 场景的参数说明、调用方式与 Ascend C 内核实现剖析【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本篇聚焦experimental/norm/layernorm_stride目录下的LayerNormStride实验算子它在标准 LayerNorm 的基础上追加了跨步stride访问能力专为 MLAMulti-head Latent Attention模型中对一行数据中前 kv_lora_rank 维做归一化、但行宽更大的张量布局而设计。读完本文你将掌握该算子的计算语义、全部参数含义与约束、基于 torch 扩展的调用方式以及其 Ascend C 内核在 AI Core 上的行级并行、UB 缓冲管理与数值计算链路。一、算子功能与计算公式根据 README 的定义算子功能MLA 中带跨步的 LayerNorm。计算公式$$ y \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum x_i^2 \epsilon}} \beta $$跨步语义在基本 LayerNorm 的基础上追加了跨步功能。输入是 shape 为[Batch, Seqlen, Stride]的张量其中Stride为跨步大小算子只对每个 Stride 中的前 kv_lora_rank 维做归一化其余维度按原样透传到输出。从公式可以看出该算子对每一行独立计算一个归一化因子1/sqrt((1/d)Σx_i² ε)再施加逐维的gammaweight缩放与betabias偏移。需要特别注意归一化窗口长度是gbW即 MLA 中的 kv_lora_rank而不是整个 Stride——这正是跨步二字的含义行宽为 Stride但有效归一化列宽为 gbW两者可以不同。在实现侧layernorm_stride.cpp 的LayerNormStrideNpu入口把 torch 张量的三个维度映射到内核参数gbH in.size(0) * in.size(1)即 Batch × Seqlen展平后的总行数gbW hiddenDimMLA 中的 kv_lora_rankstride in.size(2)即每行的实际元素数kv_lora_rank qk_rope_head_dim。见 layernorm_stride.cpp 的 launch 调用launchStatus LayerNormStrideLaunch(blockDim, stream, in.size(0) * in.size(1), hiddenDim, in.size(2), epsilon, (uint8_t*)in.data_ptr(), (uint8_t*)gamma.data_ptr(), (uint8_t*)beta.data_ptr(), (uint8_t*)out.data_ptr(), in.scalar_type() at::kHalf ? 1 : 27);二、产品支持情况产品是否支持Atlas A2 训练系列产品是源码中这一点有直接对应内核入口函数 ComputeLayernormStride 被架构宏#if __NPU_ARCH__ 2201包裹只有在该 NPU 架构上才会实例化执行模板与 README 声明的 Atlas A2 系列产品支持范围一致。三、参数说明下表完整继承自 README 的参数定义参数名输入/输出/属性描述数据类型数据格式blockDim输入AI CORE 的数量比如Ascend910B 是 40int64_t-gbW输入MLA 中的 kv_lora_rankint64_t-epsilon输入layernorm 计算的 epsfloat-in输入带跨步 layernorm 计算的输入shape 为 (Batch, Seqlen, Stride)Stride 大小为 kv_lora_rank qk_rope_head_dimBFLOAT16/HALFNDgamma输入layernorm 中的 weightBFLOAT16/HALFNDbeta输入layernorm 中的 biasBFLOAT16/HALFNDout输出等同输入BFLOAT16/HALFND结合 源码可以对各参数做进一步核对blockDim作为内核启动的 AI Core 数量传入ComputeLayernormStrideblockDim, ...入口有TORCH_CHECK(blockDim 0)校验内核启动时若gbH blockDim会自动收缩为blockDim gbH保证每个 AI Core 至多分配一行见 LayerNormStrideKernelLaunch。gbWhiddenDim归一化窗口宽度同样有hiddenDim 0校验。它决定了 UB 上临时缓冲的大小因此存在下文所述的 UB 容量上限约束。epsilonfloat 标量用于数值稳定直接参与rsqrt(均值平方和 epsilon)。in / gamma / beta / out四个张量都必须位于 NPU 设备is_npu校验且 dtype 必须是 BFLOAT16 或 HALF。gamma、beta的元素个数为 gbWout与in同形。四、调用方式README 给出的调用形式为torch.ops.ascend_ops.layernorm_stride(blockDim, gbW, epsilon, in, gamma, beta, out)该算子通过 torch extension 机制注册。在 layernorm_stride.cpp 中可以看到函数签名与 PrivateUse1即 NPU 后端实现的绑定TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m) { m.def(layernorm_stride(int blockDim, int hiddenDim, float epsilon, Tensor input, Tensor gamma, Tensor beta, Tensor output) - int); } TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m) { m.impl(layernorm_stride, LayerNormStrideNpu); }实际调用可参考仓库中的测试脚本 test_layernorm_stride.py它是一个可复制的最小完整示例import torch import torch_npu import ascend_ops_nn epsilon 1e-5 batch_size 2 seq_len 10 hidden_dim 64 dtype torch.bfloat16 input torch.randn(batch_size, seq_len, hidden_dim).to(dtype) gamma torch.randn(hidden_dim).to(dtype) beta torch.randn(hidden_dim).to(dtype) output torch.empty_like(input) input_npu input.npu() gamma_npu gamma.npu() beta_npu beta.npu() output_npu output.npu() # blockDim40 对应 Ascend910B 的 AI Core 数 torch.ops.ascend_ops_nn.layernorm_stride(40, hidden_dim, epsilon, input_npu, gamma_npu, beta_npu, output_npu)脚本中 CPU 参考实现与 NPU 结果的误差对比逻辑同样值得保留input_squared torch.square(input) mean_squared input_squared.mean(dim-1, keepdimTrue) normalized input / torch.sqrt(mean_squared epsilon) output_cpu normalized * gamma beta abs_error torch.abs(output_npu.cpu() - output_cpu) rel_error abs_error / (torch.abs(output_npu.cpu()) epsilon)测试同时遍历torch.bfloat16与torch.float16两种 dtype分别打印最大/平均绝对误差与相对误差可作为正确性验收的现成基线。注意测试脚本 import 的模块名为ascend_ops_nn即调用方使用前需要保证对应 torch extension 已构建并可通过 Python 导入。五、内核实现剖析行级 AI Core 并行与 UB 计算链路5.1 数据划分一个 AI Core 一行LayerNormStride::Init 中bkH_ 1; bkW_ gbW_; bkAlignW_ AlignUp(bkW_, 64); bkLoop_ (int64_t)(gbH_ / blockNum_); if (gbH_ % blockNum_ ! 0) { bkLoop_ 1; }可以推断出该内核的并行策略每行H 维取 1由一个 AI Core 处理gbHBatch × Seqlen按blockNum_轮转切分每个 Core 处理bkLoop_行Process()中的边界判断i * blockNum_ blockIdx_ gbH_负责处理行数不能被 Core 数整除的尾部见 Process 主循环。5.2 跨步寻址DataCopyPad 实现行内偏移拷贝in与out在 GM 中的行宽是stride_而非gbW_因此跨步体现在拷贝偏移上。CopyIn 与 CopyOut 均使用int64_t offset (process * blockNum_ blockIdx_) * stride_; DataCopyParams copyPas{1, (uint16_t)(bkW_ * sizeof(T)), 0, 0}; DataCopyPad(inLocal, inGm_[offset], copyPas, padPas);即行首按行号 × stride_定位单次只拷贝前bkW_个元素——后段qk_rope_head_dim 部分既不参与归一化计算也经由 in/out 的输入等同输出语义在调用侧保持原值。使用DataCopyPad而非普通DataCopy是为了在gbW未按硬件对齐倍数整除时由硬件完成对齐填充。5.3 计算链路均值、平方和、rsqrt 与仿射变换Compute 是核心的向量化计算段完整链路如下对齐填充与类型提升若bkAlignW_ - AlignUp(bkW_, 16) ! 0用Duplicate以 0 补齐随后把输入Cast成 float 参与后续累加规避低精度累加误差。第一轮 reduce均值Add分 64 元素一组累加到workLocalWholeReduceSum收敛出总和再乘invert_ 1.0/gbW。从源码结构看这一轮得到的均值被赋给局部变量cur_mean但最终的归一化并未做均值中心化——这与 README 公式无减均值项、仅按平方均值归一化保持一致。第二轮 reduce平方和先Mul(in, in, in)逐元素平方再经同样的分块AddWholeReduceSum求出Σx_i²乘1/d得mean(x²)随后Adds加上epsilonRsqrt得到归一化因子cur_var。仿射输出x * cur_var后依次Mul上 gamma、Add上 beta最后Cast(..., RoundMode::CAST_ROUND)转回 BF16/FP16 写出。gamma、beta 只在 Process 起始处 拷贝并升精度一次复用于本 Core 处理的全部行避免逐行重复搬运。缓冲管理上TPipe/TQue划分了 VECIN 侧的输入、gamma、beta 队列与 VECOUT 侧输出队列TBuf提供 64 float 的工作区与 16 float 的求和区成员声明。5.4 dtype 分发与 UB 容量预检dtype 分发launch 侧把at::kHalf映射为 1、at::kBFloat16映射为 27 传入内核dtype_convert.h 中的TYPE_SWITCH宏据此在half与bfloat16_t之间选择模板实例default分支回落到bfloat16_t。UB 容量预检由于每行需要占用in、out各bkAlignW_元素加上 gamma/beta 浮点工作区JudgeLayerNormStrideLaunch 在启动前按bkAlignW_ AlignUp(gbW, 64)估算 UB 用量与 184KB184 * 1024上限比较int ret JudgeLayerNormStrideLaunch(gbW, ubSize); if (ret 0) { LayerNormStrideKernelLaunch(blockDim, stream, gbH, gbW, stride, epsilon, in, gamma, beta, out, dtype); return 0; } std::cout __FUNCTION__ : UB size is limited, please check! std::endl;这是使用上的一个重要隐含约束gbWkv_lora_rank不能过大否则 UB 放不下算子会打印UB size is limited, please check!并返回失败LayerNormStrideNpu将该非零状态经int返回值向调用侧透出。5.5 执行流与 Profiling 命名Host 侧通过at_npu::native::OpCommand::RunOpApi(layerNormStride, aclCall)执行内核见 LayerNormStrideNpu在 Profiling 工具中可按layerNormStride这一名称检索到该算子的执行区间。六、约束说明与使用要点数据类型输入输出仅支持 BFLOAT16/HALFREADME 约束与 入口的四个 TORCH_CHECK 一致gamma/beta 与 in 同类型。设备要求in、gamma、beta、out 四个张量必须均位于 NPU 设备上。blockDim 与 gbW 的取值blockDim 0、hiddenDim 0blockDim 应取设备的 AI Core 数量如 Ascend910B 为 40行不足时内核会自动收缩。shape 语义in为 (Batch, Seqlen, Stride)Stride kv_lora_rank qk_rope_head_dimgbW 传 kv_lora_rankgamma/beta 长度为 gbWout 与 in 同形。gbW 上限受 UB 约束受 184KB UB 容量预检限制kv_lora_rank 过大时调用会失败并提示检查。构建集成目录 CMakeLists.txt 仅含add_sources()说明该算子源文件由上层 experimental torch extension 工程统一收集编译随 extension 包一起构建不存在独立的算子安装流程。七、相关文件索引内容路径算子说明文档参数表、公式、约束、调用说明README.mdHost 接口 Ascend C 内核实现layernorm_stride.cppdtype 宏分发half/bfloat16_tdtype_convert.hPython 正确性测试BF16/FP16 双 dtypeCPU 参考对比test_layernorm_stride.py构建文件CMakeLists.txt综上LayerNormStride 以行级 AI Core 并行 跨步 DataCopyPad UB 内两级 reduce的轻量结构实现了 MLA 场景下对宽行张量前 kv_lora_rank 维的归一化。使用时需牢记 BF16/FP16 限定、stride 与 gbW 的 shape 约定以及 gbW 过大会触发 UB 容量检查失败这三点约束。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表