ARTICLE DETAIL

资讯详情

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

ccv 中融合 Dropout、残差与 LayerNorm 的 CUDA 算子源码剖析:基于 FlashAttention layer_norm 扩展

ccv 中融合 Dropout、残差与 LayerNorm 的 CUDA 算子源码剖析:基于 FlashAttention layer_norm 扩展 计算机视觉深度学习【免费下载链接】ccvC-based/Cached/Core Computer Vision Library, A Modern Computer Vision Library项目地址https://gitcode.com/gh_mirrors/cc/ccv点击查看免费下载本文聚焦 ccv 仓库内 vendored 的 FlashAttention layer_norm CUDA 扩展README讲解它如何在单次 kernel launch 中融合 dropout、残差相加与 LayerNorm并可选 RMSNorm、parallel residual以及在 ccv 的神经网络计算图NNC框架中如何被复用为 GPU 上的高效 LayerNorm 实现。读完本文你将掌握该扩展的隐藏维度特化机制、五元类型系统、编译期开关模板以及它在 ccv CUDA 兼容层中的集成与调用方式。一、背景为什么要融合 Dropout、残差与 LayerNormTransformer 的 Pre-Norm 与 Post-Norm 架构中残差连接、dropout 与归一化层是高频出现的最小算子组合。若三个算子各自独立 launch kernel不仅要多轮读写全局显存还会引入额外的同步开销。原文档指出这个 CUDA 扩展的核心目标即是在单个 kernel内完成dropout按行生成随机 mask 并缩放残差相加x0 residualLayerNorm 归一化计算均值/方差做(x - mu) * rs * gamma beta。该实现基于 NVIDIA Apex 的 FastLayerNorm并在其基础上做了几处关键改动对应 README新增 dropout 与残差融合路径同时支持 pre-norm 与 post-norm 两种架构支持更多隐藏维度所有能被 8 整除的维度上限 8192提供 RMSNorm 选项省略均值项支持 parallel residualGPT-J、GPT-NeoX、PaLM 一类架构两条分支分别归一化。二、源码布局一个独立的 CUDA 算子子模块该扩展在仓库中位于 lib/nnc/gpu/3rdparty/flash_attn/layer_norm/共 22 个文件职责划分清晰文件职责ln.h参数结构ParamsBase/FwdParams/BwdParams、类型 ID 与注册器、LaunchParamsln_api.h / ln_api.cu对外 APIDataType枚举、round_hidden_size、run_layer_norm_fwdln_fwd_kernels.cuh前向 kernel 本体与launch_模板ln_kernel_traits.h编译期 kernel 特征线程数、每次全局加载字节数、CTA 排布ln_utils.cuh向量化加载Vec、warp/CTA 归约Reducer/Stats、CTA 间同步InterCTASyncstatic_switch.hBOOL_SWITCH编译期布尔分派宏ln_fwd_256.cu ~ ln_fwd_8192.cu14 个隐藏维度的特化注册文件256、512、768、1024、1280、1536、2048、2560、3072、4096、5120、6144、7168、8192注意该目录中的ln_api.cu、ln_fwd_*.cu会被 ccv 的 CUDA 编译系统直接纳入构建见 lib/nnc/gpu/makefile而不是像原 FlashAttention 仓库那样通过pip install单独安装。三、核心 API 与隐藏维度的对齐策略3.1 数据类型枚举ln_api.h 定义了三种基础数据类型enum DataType { DATA_TYPE_FP16 0, DATA_TYPE_BF16 1, DATA_TYPE_FP32 2, };run_layer_norm_fwd的完整签名会同时指定权重类型、输入类型、残差类型、输出类型与隐藏维度计算类型固定为 FP32bool run_layer_norm_fwd(LaunchParamsFwdParams launch_params, DataType wtype, DataType itype, DataType rtype, DataType otype, uint32_t hidden_size, bool configure_params);3.2 round_hidden_size非对齐维度的向上取整ln_api.cu 中实现了维度对齐uint32_t round_hidden_size(uint32_t hidden_size) { const uint32_t multiple hidden_size 1536 ? 256 : (hidden_size 3072 ? 512 : 1024); return (hidden_size multiple - 1) / multiple * multiple; }即 1536按 256 的倍数对齐 3072按 512 对齐否则按 1024 对齐。kernel 分派时先对实际hidden_size取整再查注册表见 ln_api.cu。这也是原文档「所有能被 8 整除、最大 8192」约束的工程化体现——kernel 内通过Is_even_cols分支处理实际列数与特化维度不一致的尾部元素。3.3 五元类型组合 → 64 位函数键ln.h 用「权重 / 输入 / 残差 / 输出 / 计算」五种类型的组合生成函数键每种类型用 2 位编码fp160、bf161、fp322依次左移 0/2/4/6/8 位再左移 32 位后与hidden_size合并成 64 位键FwdRegistrar/BwdRegistrar等注册器在静态初始化阶段把实例化的 kernel 函数写入FWD_FUNCS/BWD_FUNCS等全局unordered_mapln.h。以 ln_fwd_4096.cu 为例每个隐藏维度文件通过宏REGISTER_FWD_LAUNCHER注册 10 种类型组合fp32/fp16/bf16 的权重与输入、残差、输出的多种搭配计算类型统一为 fp32例如REGISTER_FWD_LAUNCHER( 4096, fp16, fp16, fp16, fp16, fp32, 1, 1, 4, 16); REGISTER_FWD_LAUNCHER( 4096, bf16, bf16, bf16, bf16, fp32, 1, 1, 4, 16);宏的末 4 个整数参数依次是CTAS_PER_ROW每行 CTA 数、WARPS_M行方向 warp 数、WARPS_N列方向 warp 数、BYTES_PER_LDG每次全局加载字节数通常 16即 16 字节向量化读写。四、前向 kernel 的原理单 kernel 内完成全部融合4.1 编译期开关消除运行时分支ln_fwd_kernels.cuh 中的 kernel 模板带 5 个编译期 bool 参数templatetypename Ktraits, bool Is_dropout, bool Has_colscale, bool Has_subset, bool Is_even_cols __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_kernel(FwdParams params)而 launch_ 通过 static_switch.h 的BOOL_SWITCH宏在运行时把 4 个布尔条件展开成 16 个完全特化的 kernel 实例避免热路径上的分支判断BOOL_SWITCH(launch_params.params.dropout_keep_p 1.f, IsDropoutConst, [] { BOOL_SWITCH(has_colscale, HasColscaleConst, [] { BOOL_SWITCH(has_subset, HasSubsetConst, [] { BOOL_SWITCH(is_even_cols, IsEvenColsConst, [] { auto kernel ln_fwd_kernelKernel_traits, IsDropoutConst, HasColscaleConst, HasSubsetConst, IsEvenColsConst; ... }); }); }); });4.2 单行计算流程从 ln_fwd_kernels.cuh 可以还原每一行数据的处理链路读入 gamma/beta/colscale每个线程按向量化索引预取归一化权重gamma、偏置beta可空则置零、可选列缩放colscale生成 dropout mask若启用 dropout用curand_init初始化 Philox 状态ln_fwd_kernels.cuh逐元素curand_uniform(state) dropout_keep_p决定保留与否并乘上dropout_scale即 1/keep_p 的逆缩放融合残差x_ij has_residual ? x0_ij residual_ij : x0_ij若启用rowscale或colscale会先对x0逐行/逐列缩放统计量计算Stats::compute汇总均值mu与二阶矩m2利用inverse_cols1/cols归一化归一化与仿射compute_t rs rsqrtf(m2 * params.inverse_cols params.epsilon (!params.is_rms_norm ? 0.f : mu * mu)); compute_t y_ij compute_t(rs * (xf[it * NUM_ELTS jt] - (!params.is_rms_norm ? mu : 0.f))); z.data.elt[jt] output_t(g_ij * y_ij b_ij);注意这里 RMSNorm 与 LayerNorm 共用一个rs公式RMS 模式下额外加上mu*mu即均方并在减均值时跳过mu项LayerNorm 模式下m2*inverse_cols即方差。均值mu与逆标准差rs会写入mu_ptr[row]/rs_ptr[row]saved stats供反向传播复用。4.3 归约与 CTA 间协作ln_utils.cuh 提供了分层归约基础设施Vec按BYTES_PER_LDG默认 16 字节把连续元素打包成uint4等宽类型做向量化 load/storeln_utils.cuhReducercompute_t, CTAS_PER_ROW, WARPS_M, WARPS_N先 warp 内__shfl_down_sync/__shfl_xor_sync归约再经 shared memory 合并多个 warp当CTAS_PER_ROW 1时通过全局 workspace 与InterCTASync栅栏做 CTA 间归约ln_utils.cuhStats用{m, m2}均值、二阶矩两元组交换而非三元组warp_chan_upd_dynamic实现多 batch 的在线合并ln_utils.cuh。前向 kernel 目前在static_assert(CTAS_PER_ROW 1)约束下工作ln_fwd_kernels.cuh因此归约主要发生在 warp 与 CTA 内部launch_中CTAS_PER_ROW 1分支保留的 workspace/barrier 计算ln_fwd_kernels.cuh是为未来多 CTA 协同预留的路径。4.4 配置与 launchconfigure_params true时launch_ 会用cudaOccupancyMaxActiveBlocksPerMultiprocessor计算每 SM 可驻留 CTA 数据此决定ctas_per_col网格规模、elts_per_thread、workspace 与 barrier 大小正式 launch 时若动态共享内存超过 48KB 则通过cudaFuncSetAttribute提高上限ln_fwd_kernels.cuh。五、ccv 集成如何把该算子接入 NNC 计算图在原 FlashAttention 仓库中该扩展已自 2024-01-05 起被 Triton 实现替换不再默认使用见 README但在 ccv 中它仍以 vendored 源码形式服务于 GPU 后端构建接入lib/nnc/gpu/makefile 把ln_api.cu与 14 个ln_fwd_*.cu全部纳入FLASH_ATTN_SRCS与 FlashAttention 主算子一同编译并链接进libnnc-compat-cuda-sm80.osm80 即 Ampere 架构与 README「仅在 A100 上测试过」的算力要求一致C 层封装ccv_nnc_flash_norm_gpu.h 是 ccv 侧的中转层仅当HAVE_CUDNN HAVE_CUDA_SM80时编译它把 ccv 张量类型CCV_16F/CCV_16BF/CCV_32F映射到layer_norm::DataTypeccv_nnc_flash_norm_gpu.h并在_ccv_nnc_flash_norm_check中强制约束张量必须连续、归一化轴必须是尾轴、cols % 8 0且round_hidden_size(cols) 8192ccv_nnc_flash_norm_gpu.hlaunch 配置缓存_ccv_nnc_flash_norm_configure以(device, rows, cols, input_datatype, weight_datatype)为键缓存 16 个槽位的占用率配置避免重复查询cudaOccupancyMaxActiveBlocksPerMultiprocessorccv_nnc_flash_norm_gpu.h命令注册归一化命令CCV_NNC_LAYER_NORM_FORWARDccv_nnc_cmd.h在 GPU/CUDNN 后端注册ccv_nnc_cmd.inc并可通过is_rms_norm开关复用同一套 kernel 完成 RMSNormccv_nnc_flash_norm_gpu.h 中params.is_rms_norm is_rms_norm。ccv 侧调用时不会启用 dropoutdropout_keep_p 1.f见 ccv_nnc_flash_norm_gpu.h即该算子最完整的融合形态dropout residual norm保留在原始 FlashAttention 使用场景中而 ccv 的 NNC 框架复用了它的高性能归一化与统计量计算路径。六、使用约束与注意事项综合原文档与源码使用该扩展前需确认以下前提隐藏维度必须是 8 的倍数round_hidden_size(hidden_size) 8192超过 8k 的场景原文档建议 filed issue仓库侧在 ccv_nnc_flash_norm_gpu.h 直接拒绝硬件原文档声明仅在 A100 上测试ccv 构建时以HAVE_CUDA_SM80为编译开关意味着面向 Amperesm80及以上算力数据类型权重/输入/残差/输出支持 fp16、bf16、fp32 的常见组合计算统一在 fp32 中进行以保证数值稳定归一化轴ccv 集成要求归一化轴必须是张量的尾部连续轴_ccv_nnc_flash_norm_tail_axes校验见 ccv_nnc_flash_norm_gpu.h安装方式差异原 README 中的cd csrc/layer_norm pip install .针对 FlashAttention 仓库的独立 Python 扩展在 ccv 中无需单独安装编译 ccv 的 CUDA 兼容层即可源码已由 lib/nnc/gpu/makefile 纳入构建。七、小结这个 layer_norm 子模块的价值在于把「dropout 残差 归一化」这条 Transformer 高频路径压缩进单个 CUDA kernel通过隐藏维度特化2568192 共 14 档、五元类型键值注册表、BOOL_SWITCH编译期开关以及 warp/CTA 两级归约与 16 字节向量化访存同时覆盖 LayerNorm 与 RMSNorm、pre-norm 与 post-norm、parallel residual 等多种架构形态。在 ccv 中它作为 FlashAttention 第三方源码的一部分被编译进 sm80 CUDA 兼容层并经 ccv_nnc_flash_norm_gpu.h 接入CCV_NNC_LAYER_NORM_FORWARD命令为 GPU 上的归一化算子提供了融合实现参考。若需在更大隐藏维度或非 Ampere 硬件上使用需按源码中的注册与开关机制自行扩展特化文件。赞分享计算机视觉深度学习【免费下载链接】ccvC-based/Cached/Core Computer Vision Library, A Modern Computer Vision Library项目地址https://gitcode.com/gh_mirrors/cc/ccv点击查看免费下载相关推荐CANN ops-transformer 融合算子 LayerNormAddAddQuant 深度解析残差加法、LayerNorm 与动态量化的三合一 Kernel 设计CANN ops transformer 融合算子 LayerNormAddAddQuant 深度解析残差加法、LayerNorm 与动态量化的三合一 Ker算子库人工智能大模型深度学习CANNAscendCANN ops-nn AdaLayerNorm 算子深度解析自适应 LayerNorm 融合算子的原理、API 调用与源码实现CANN ops nn AdaLayerNorm 算子深度解析自适应 LayerNorm 融合算子的原理、API 调用与源码实现 导读 AdaLayerNor人工智能算子库深度学习CANNAscendCANN ops-transformer 中的 BlockAttentionResidualsGrad融合 Softmax 与 RMSNorm 反向的注意力残差梯度算子解析CANN ops transformer 中的 BlockAttentionResidualsGrad融合 Softmax 与 RMSNorm 反向的注意力残算子库人工智能大模型深度学习CANNAscend上一篇Typhoeus与Faraday集成完全教程构建企业级HTTP请求系统下一篇Munder Difflin 实战用 GitHub Webhook 把仓库事件变成 AI Agent 任务创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表