CUTLASS 111_hopper_ssd 实战基于 Hopper 架构的 SSD状态空间分解CUDA 实现解析【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass导读本文围绕 CUTLASS 仓库中的examples/111_hopper_ssd示例讲解如何在 NVIDIA Hopper GPUcompute capability 9.0上以 CUTLASS 3.x 组件实现 SSDState Space Decomposition状态空间分解算子。SSD 是状态空间模型SSM类架构中把序列变换分解为块内计算与跨块状态传播的高效数学形式该示例给出了从命令行驱动、设备端 API、warp 特化内核、TMA 流水线到参考实现验证的完整工程化方案。读完本文你将掌握该示例的构建方法、全部命令行参数含义、源码目录结构、四阶段 BMMIntraBMM1/IntraBMM2/InterBMM1/InterBMM2流水线设计以及其当前支持范围与性能特征。一、示例概述SSD 在 CUTLASS 中的落地形态原文档examples/111_hopper_ssd/README.md明确说明本示例演示在 NVIDIA Hopper 架构上实现 SSD 操作核心是利用 CUTLASS 库组件做高性能张量计算充分发挥 Hopper 的 TMATensor Memory Accelerator与 warp specialization 能力。从源码结构看该示例并非简单的 GEMM 调用而是把 SSD 的数学流程拆解为多个可以在 Hopper 上并行执行的矩阵乘阶段主入口 111_hopper_ssd.cu命令行解析、数据初始化、内核运行与结果验证设备层封装 device/ssd.hpp提供cutlass::ssd::device::SSD这一符合 CUTLASS 3.x 约定的设备算子 API内核层 kernel/sm90_ssd_kernel_tma_warpspecialized.hppwarp 特化的持久化内核主体集体层 collective/sm90_ssd_gemm_tma_warpspecialized.hpp 与 collective/sm90_ssd_epilogue.hpp主循环与 epilogue 的具体实现参考实现 reference/reference_ssd.hpp 与 reference/reference_ssd_cumsum.hpp用于数值验证的朴素实现。在 reference/reference_ssd.hpp 中定义了PHASE宏#define PHASE 0注释标注 Phase 0 为训练training、Phase 1 为推理inference说明该参考实现目前面向训练场景。二、系统要求原文档给出的硬性前置条件如下这与主程序 111_hopper_ssd.cu 中的运行时检查一一对应要求说明NVIDIA GPUHopper 架构compute capability 9.0SM90CUDA Toolkit12.0 或更新版本编译器支持 C17 的编译器在main()中程序会查询设备属性并做如下判定若__CUDACC_VER_MAJOR__ 12 || props.major 9则提示需要 Hopper 及以上架构、CUDA 12.0 及以上并直接返回若 compute capability 不是 9.0props.major ! 9 || props.minor ! 0则提示需要严格为 Hoppercompute capability 90的 GPU。同时整个示例主体被包裹在#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)条件编译中只有开启了 SM90 MMA 支持即用支持 SM90 的 CUDA 编译器构建时才会编译执行。运行示例前请用nvidia-smi确认 GPU 型号为 H100/H200 等 SM90 设备。三、构建示例原文档指出Follow the cutlass example building遵循 CUTLASS 示例的通用构建方式。构建目标由 examples/111_hopper_ssd/CMakeLists.txt 定义通过cutlass_example_add_executable(111_hopper_ssd 111_hopper_ssd.cu)生成名为111_hopper_ssd的可执行文件并单独为该源文件附加--use_fast_math编译选项——这与内核中大量使用expf()等超越函数直接相关快速数学模式有助于提升这类 ALU 密集算子的吞吐。典型构建流程在仓库根目录执行# 配置显式指定 GPU 架构与示例开关 cmake -S . -B build -DCUTLASS_ENABLE_EXAMPLESON -DCUTLASS_NVCC_ARCHS90a # 只编译本示例 cmake --build build --target 111_hopper_ssd -j # 运行需在 SM90 GPU 上 ./build/examples/111_hopper_ssd/111_hopper_ssd-DCUTLASS_NVCC_ARCHS90a指定只生成 SM90 架构代码如你的环境已有现成构建目录也可直接复用其中的示例构建产物。四、命令行选项详解原文档列出了示例支持的 6 个命令行选项。结合 111_hopper_ssd.cu 中Options结构体的解析逻辑与print_usage()输出整理出完整语义选项类型默认值说明--help标志关打印使用说明后退出--iterationsint整数1基准测试迭代次数--without_verify标志关跳过结果验证--verbose标志关打印每个内核的执行时间--Gint整数2Group 数--Bint整数3Batch 数--Eint整数2扩展因子Expanded factor--Hint整数2头数Number of heads需要注意源码中的几个联动行为EH E * H扩展因子 × 头数问题形状problem shape被定义为七元组(G, B, EH, C, L, D, N)运行时会在启动前打印例如默认参数下为(2, 3, 4, 8, 128, 64, 128)当--iterations 1时代码会自动把measure与verbose置为true进入计时与逐核打印模式静态尺寸C 8、D 64、L 128、N 128被硬编码为Int常量参考内核暂不支持动态CG 的实际支持范围命令行默认G 2但 111_hopper_ssd.cu 的initialize()中有assert(g 1 Only group size 1 is supported)。也就是说当前实现实际上仅接受--G1否则会在初始化阶段触发断言失败这是源码中明确可验证的限制在--verbose模式下程序会分别打印 cumsum 内核与 SSD 内核的平均运行时间runtime_ms并输出共享内存用量[Usage] smem与smem size两类打印前者来自 device/ssd.hpp后者来自主程序。运行示例# 单次运行并验证结果 ./111_hopper_ssd --G1 --B3 --E2 --H2 # 基准测试100 次迭代自动开启 warmup3 与 verbose ./111_hopper_ssd --G1 --iterations100 # 跳过验证仅测速 ./111_hopper_ssd --G1 --without_verify --verbose五、SSD 的数学流程与参考实现要读懂内核先理解参考实现定义的计算。在 reference/reference_ssd.hpp 中可以看到两个核心朴素算子1.segsum分段求和对输入按最后一维做前缀和得到cum_sum随后构造一个下三角指数衰减矩阵——对角线为 1下三角第(i, j)项为expf(cum_sum(i) - cum_sum(j))上三角为 0。这一矩阵正是 SSD 中块内状态传播权重也是原文档 Performance 部分所说性能受限于 SEGSUM 计算部分的对象。2.mma通用矩阵乘朴素三重循环实现用于生成与内核结果比对的参考输出。参考路径还包括 reference/reference_ssd_cumsum.hpp 中的CumsumKernel它在主程序初始化阶段于设备端对DeltaA张量预先执行 cumsum111_hopper_ssd.cu并把结果tensor_DeltaA_cumsum作为 SSD 内核的输入之一。验证逻辑位于TestBed::run()与compare_reference验证开启时调用ssd_reference生成参考Y与F张量再与内核输出逐元素比较默认容差epsilon 0.05f同时采用绝对误差与相对误差abs_error / (max(|act|, |ref|) 1e-5)两种判据任一误差超过阈值即判定失败并打印首个不匹配位置及参考/计算值矩阵。验证通过输出everything is ok.失败输出something is wrong!!!!!。六、源码结构从 Builder 到设备 API该示例的模板分层体现了 CUTLASS 3.x 的Builder → Kernel → Device组装模式kernel/sm90_ssd_kernel_builder.hppSm90SsdBuilder承担编译期装配。它接收模板参数Element/ElementDA/ElementAcc/ElementY、TileShape即 (L, D, N)以及HAS_D/D_HAS_HDIM/HAS_Z三个布尔开关并固定若干关键配置StagesX StagesY 2、StagesZ 1注释明确说明这是共享内存容量限制所致、epilogue 采用TmaWarpSpecialized调度、partial_m 128。最终组合出CollectiveMainloop、CollectiveEpilogue与PersistentTileScheduler产出SsdKernelTmaWarpSpecialized内核。kernel/sm90_ssd_tile_scheduler.hppPersistentTileScheduler实现持久化网格num_blocks B * EH网格大小取min(num_blocks, sm_count)每个线程块通过block_idx gridDim.x步进遍历多个 tile从而在较少 block 下充分利用全部 SM。它同时负责把block_idx解码为(b, eh)与 group 坐标供不同 producer warp 使用。kernel/sm90_ssd_kernel_tma_warpspecialized.hpp内核主体SsdKernelTmaWarpSpecialized其中NumMmaWarpGroups 2硬编码NumLoadWarpGroups 1即每个线程块 3 个 warp groupMaxThreadsPerBlock 3 * NumThreadsPerWarpGroup定义三种WarpGroupRoleProducer、Consumer0、Consumer1Producer 内部再按 lane 划分四种ProducerWarpRoleLoadX、LoadDelta、LoadBC、LoadZ为 X、Delta/DeltaA、B、C、D、Z 以及协作与存储建立了多条 TMA/pipelinePipelineTmaAsync/PipelineTmaStore/PipelineAsync通过warpgroup_reg_alloc/dealloc做寄存器重分配LoadRegisterRequirement 40 - 2*8体现 Hopper 上 producer/consumer 分时复用寄存器的优化手段。device/ssd.hppcutlass::ssd::device::SSDKernel是标准的 CUTLASS 3.x 设备算子门面提供can_implement、get_workspace_size、initialize、update、run与operator()等接口当ArchTag::kMinComputeCapability 90时通过ClusterLauncher发起带 cluster 形状的扩展启动。主程序 111_hopper_ssd.cu 的调用顺序为构造SsdOperation::Arguments含输入输出张量指针与经layout*_transformed()变换的布局→get_workspace_size分配 workspace →can_implement校验 →initialize初始化参数 → warmup默认 3 次→ 正式迭代计时 → 参考实现验证。其中KernelHardwareInfo通过query_device_multiprocessor_count查询 SM 数并传给 tile scheduler。七、主循环四个 BMM 阶段的流水线设计collective/sm90_ssd_gemm_tma_warpspecialized.hpp 中的SsdMainloopTmaWarpSpecialized把 SSD 计算组织为四个矩阵乘阶段每个阶段均复用 CUTLASSCollectiveBuilder生成KernelTmaWarpSpecialized类型的集体算子阶段TileShape布局说明IntraBMM1(L, L, N)NT块内第一阶段矩阵乘对应 B 与 C 张量IntraBMM2(L, D, N)TN块内第二阶段矩阵乘对应累积结果与 XInterBMM1(N, D, L)TN跨块第一阶段对应 B 与 XInterBMM2(L, D, N)NN跨块第二阶段对应状态 P 与 CTileShape (L, D, N) (128, 64, 128)ClusterShape被强制为1x1x1不支持多播。数据通路方面X、B、C 使用 TMA 加载SM90_TMA_LOADAlignment 16/sizeof(Element)保证 16 字节对齐Delta/DeltaA 则使用SM90_BULK_COPY_AUTO批量拷贝。从内核的 Consumer 侧代码可以还原数据流Consumer0Intra 路径等待 B、C →mma_intra_1做 IntraBMM1 →pre_intra_2读取 Delta/DeltaA转换为float后按m n的下三角条件计算expf(deltaA_col - deltaA_row)与 Delta 及 IntraBMM1 结果逐元素相乘即 SEGSUM 的逐元素实现→mma_intra_2用 X 做 IntraBMM2 →store_intra通过协作 pipeline 把部分和交给 epilogue。Consumer1Inter 路径state_init初始化状态张量P含一次经共享内存的转置配合TransposeBarrier/TransformBarrier命名屏障→pre_inter_1等待 B 与 Delta/DeltaA用SM75_U32x4_LDSM_N从共享内存装载 B 片段利用__shfl_sync广播每行最后一个 DeltaA 列值计算expf(last_column - deltaA) * delta * b的加权 B →mma_inter_1用 X 做 InterBMM1 →pre_inter_2把上一轮状态按expf(last_column)衰减后累加进结果并更新状态 →mma_inter_2用状态 P 与 C 做 InterBMM2同时对 DeltaA 取expf供 epilogue 使用 →post_inter_2把新状态经转置写回共享内存。代码注释sm90_ssd_gemm_tma_warpspecialized.hpp特别指出Inter 路径的 DeltaA 处理依赖 GMMA 已知的 TV 分区方式仅对 128x64x128 tile 成立布局为((REG_N,REG_M,REG_N_REP),ATOM_M_REP,ATOM_N_REP)并有static_assert强制GMMA::ALayout_64x16——这正是原文档Only support LxDxN 128x64x128限制的底层原因。八、Epilogue输出张量、状态与前向最终态collective/sm90_ssd_epilogue.hpp 中的SsdEpilogue模板类负责收尾工作维护三条 TMA 流水线EpiloadPipelineD装载对角张量 D、EpiloadPipelineZ装载额外张量 Z、StorePipeline/StorePPipelineTMA 存储 Y 与最终状态 F布尔模板开关HAS_D、D_HAS_HDIM、HAS_Z分别控制是否计算 D、D 是否含 head 维度、是否输出 Z对应共享内存布局与kEpiloadDBytes/kEpiloadZBytes的编译期裁剪输出张量规格在 111_hopper_ssd.cu 中有明确注释x [b, eh, d, c, l]、delta/delta_A [b, eh, c, l]、B/C [b, g, n, c, l]、y [b, eh, d, c, l]、fstate [b, eh, d, n]。epilogue 的store阶段会汇总 Consumer1 产出的tInter2、tDeltaA与更新后的tD按 chunk 写出 Y全部 chunk 完成后由store_p写出最终状态 F。若D_HAS_HDIM为真还需在结尾对 D 流水线做一次consumer_release并推进状态。九、当前限制与性能特征原文档明确列出两项限制结合源码可以给出更精确的边界仅支持 LxDxN 128x64x128TileShape在 111_hopper_ssd.cu 固定为(Options::L, Options::D, Options::N)而 L/D/N 是编译期常量GMMA 布局相关的static_assert进一步把 Inter 路径限定在 128 宽的 tile 上。此外partial_m 128也是写死的。输入输出使用 bfloat16 精度Options::Element cutlass::bfloat16_t累加器为floatDeltaA 张量用float存储而 Delta/B/C/D/Z 均为 bf16。参考实现中的逐元素运算expf等也全部提升到float后计算再回写。性能特征方面原文档给出的三点均有源码佐证利用 TMAX/B/C 走SM90_TMA_LOAD描述符并在启动阶段由单个线程prefetch_tma_descriptors预取描述符warp 特化Producer 按 lane 分工同时推进 X、Delta/DeltaA、B/C、D/Z 的装载两个 Consumer warp group 分别执行 Intra 与 Inter 路径配合PipelineTmaAsync实现多级流水受限于 SEGSUM、ALU boundSEGSUM 涉及共享内存中 Delta/DeltaA 的读写、expf指数运算与逐元素乘加见pre_intra_2、pre_inter_1这些标量浮点操作不经过张量核构成性能瓶颈--use_fast_math编译选项正是为缓解这一 ALU 密集特性。需要注意的是原文档未给出任何具体性能数值本文同样不做任何性能数据推断。十、快速自查清单硬件/工具链确认 GPU 为 SM90compute capability 9.0、CUDA ≥ 12.0、编译器支持 C17构建开启CUTLASS_ENABLE_EXAMPLES并指定CUTLASS_NVCC_ARCHS90a目标111_hopper_ssd运行当前实现仅接受--G1--iterations1会自动进入计时模式并打印 cumsum/ssd 内核耗时与共享内存占用验证默认执行设备端参考实现对比容差 0.05可用--without_verify关闭边界tile 固定为 128x64x128精度为 bf16 输入 / fp32 累加。版权说明本示例源码与文档版权归 NVIDIA CORPORATION AFFILIATES采用 BSD-3-Clause 许可证SPDX-License-Identifier: BSD-3-Clause具体条款见 README.md 与各源文件头部使用与分发时请遵守该许可证要求。【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CUTLASS 111_hopper_ssd 实战：基于 Hopper 架构的 SSD（状态空间分解）CUDA 实现解析
CUTLASS 111_hopper_ssd 实战基于 Hopper 架构的 SSD状态空间分解CUDA 实现解析【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass导读本文围绕 CUTLASS 仓库中的examples/111_hopper_ssd示例讲解如何在 NVIDIA Hopper GPUcompute capability 9.0上以 CUTLASS 3.x 组件实现 SSDState Space Decomposition状态空间分解算子。SSD 是状态空间模型SSM类架构中把序列变换分解为块内计算与跨块状态传播的高效数学形式该示例给出了从命令行驱动、设备端 API、warp 特化内核、TMA 流水线到参考实现验证的完整工程化方案。读完本文你将掌握该示例的构建方法、全部命令行参数含义、源码目录结构、四阶段 BMMIntraBMM1/IntraBMM2/InterBMM1/InterBMM2流水线设计以及其当前支持范围与性能特征。一、示例概述SSD 在 CUTLASS 中的落地形态原文档examples/111_hopper_ssd/README.md明确说明本示例演示在 NVIDIA Hopper 架构上实现 SSD 操作核心是利用 CUTLASS 库组件做高性能张量计算充分发挥 Hopper 的 TMATensor Memory Accelerator与 warp specialization 能力。从源码结构看该示例并非简单的 GEMM 调用而是把 SSD 的数学流程拆解为多个可以在 Hopper 上并行执行的矩阵乘阶段主入口 111_hopper_ssd.cu命令行解析、数据初始化、内核运行与结果验证设备层封装 device/ssd.hpp提供cutlass::ssd::device::SSD这一符合 CUTLASS 3.x 约定的设备算子 API内核层 kernel/sm90_ssd_kernel_tma_warpspecialized.hppwarp 特化的持久化内核主体集体层 collective/sm90_ssd_gemm_tma_warpspecialized.hpp 与 collective/sm90_ssd_epilogue.hpp主循环与 epilogue 的具体实现参考实现 reference/reference_ssd.hpp 与 reference/reference_ssd_cumsum.hpp用于数值验证的朴素实现。在 reference/reference_ssd.hpp 中定义了PHASE宏#define PHASE 0注释标注 Phase 0 为训练training、Phase 1 为推理inference说明该参考实现目前面向训练场景。二、系统要求原文档给出的硬性前置条件如下这与主程序 111_hopper_ssd.cu 中的运行时检查一一对应要求说明NVIDIA GPUHopper 架构compute capability 9.0SM90CUDA Toolkit12.0 或更新版本编译器支持 C17 的编译器在main()中程序会查询设备属性并做如下判定若__CUDACC_VER_MAJOR__ 12 || props.major 9则提示需要 Hopper 及以上架构、CUDA 12.0 及以上并直接返回若 compute capability 不是 9.0props.major ! 9 || props.minor ! 0则提示需要严格为 Hoppercompute capability 90的 GPU。同时整个示例主体被包裹在#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)条件编译中只有开启了 SM90 MMA 支持即用支持 SM90 的 CUDA 编译器构建时才会编译执行。运行示例前请用nvidia-smi确认 GPU 型号为 H100/H200 等 SM90 设备。三、构建示例原文档指出Follow the cutlass example building遵循 CUTLASS 示例的通用构建方式。构建目标由 examples/111_hopper_ssd/CMakeLists.txt 定义通过cutlass_example_add_executable(111_hopper_ssd 111_hopper_ssd.cu)生成名为111_hopper_ssd的可执行文件并单独为该源文件附加--use_fast_math编译选项——这与内核中大量使用expf()等超越函数直接相关快速数学模式有助于提升这类 ALU 密集算子的吞吐。典型构建流程在仓库根目录执行# 配置显式指定 GPU 架构与示例开关 cmake -S . -B build -DCUTLASS_ENABLE_EXAMPLESON -DCUTLASS_NVCC_ARCHS90a # 只编译本示例 cmake --build build --target 111_hopper_ssd -j # 运行需在 SM90 GPU 上 ./build/examples/111_hopper_ssd/111_hopper_ssd-DCUTLASS_NVCC_ARCHS90a指定只生成 SM90 架构代码如你的环境已有现成构建目录也可直接复用其中的示例构建产物。四、命令行选项详解原文档列出了示例支持的 6 个命令行选项。结合 111_hopper_ssd.cu 中Options结构体的解析逻辑与print_usage()输出整理出完整语义选项类型默认值说明--help标志关打印使用说明后退出--iterationsint整数1基准测试迭代次数--without_verify标志关跳过结果验证--verbose标志关打印每个内核的执行时间--Gint整数2Group 数--Bint整数3Batch 数--Eint整数2扩展因子Expanded factor--Hint整数2头数Number of heads需要注意源码中的几个联动行为EH E * H扩展因子 × 头数问题形状problem shape被定义为七元组(G, B, EH, C, L, D, N)运行时会在启动前打印例如默认参数下为(2, 3, 4, 8, 128, 64, 128)当--iterations 1时代码会自动把measure与verbose置为true进入计时与逐核打印模式静态尺寸C 8、D 64、L 128、N 128被硬编码为Int常量参考内核暂不支持动态CG 的实际支持范围命令行默认G 2但 111_hopper_ssd.cu 的initialize()中有assert(g 1 Only group size 1 is supported)。也就是说当前实现实际上仅接受--G1否则会在初始化阶段触发断言失败这是源码中明确可验证的限制在--verbose模式下程序会分别打印 cumsum 内核与 SSD 内核的平均运行时间runtime_ms并输出共享内存用量[Usage] smem与smem size两类打印前者来自 device/ssd.hpp后者来自主程序。运行示例# 单次运行并验证结果 ./111_hopper_ssd --G1 --B3 --E2 --H2 # 基准测试100 次迭代自动开启 warmup3 与 verbose ./111_hopper_ssd --G1 --iterations100 # 跳过验证仅测速 ./111_hopper_ssd --G1 --without_verify --verbose五、SSD 的数学流程与参考实现要读懂内核先理解参考实现定义的计算。在 reference/reference_ssd.hpp 中可以看到两个核心朴素算子1.segsum分段求和对输入按最后一维做前缀和得到cum_sum随后构造一个下三角指数衰减矩阵——对角线为 1下三角第(i, j)项为expf(cum_sum(i) - cum_sum(j))上三角为 0。这一矩阵正是 SSD 中块内状态传播权重也是原文档 Performance 部分所说性能受限于 SEGSUM 计算部分的对象。2.mma通用矩阵乘朴素三重循环实现用于生成与内核结果比对的参考输出。参考路径还包括 reference/reference_ssd_cumsum.hpp 中的CumsumKernel它在主程序初始化阶段于设备端对DeltaA张量预先执行 cumsum111_hopper_ssd.cu并把结果tensor_DeltaA_cumsum作为 SSD 内核的输入之一。验证逻辑位于TestBed::run()与compare_reference验证开启时调用ssd_reference生成参考Y与F张量再与内核输出逐元素比较默认容差epsilon 0.05f同时采用绝对误差与相对误差abs_error / (max(|act|, |ref|) 1e-5)两种判据任一误差超过阈值即判定失败并打印首个不匹配位置及参考/计算值矩阵。验证通过输出everything is ok.失败输出something is wrong!!!!!。六、源码结构从 Builder 到设备 API该示例的模板分层体现了 CUTLASS 3.x 的Builder → Kernel → Device组装模式kernel/sm90_ssd_kernel_builder.hppSm90SsdBuilder承担编译期装配。它接收模板参数Element/ElementDA/ElementAcc/ElementY、TileShape即 (L, D, N)以及HAS_D/D_HAS_HDIM/HAS_Z三个布尔开关并固定若干关键配置StagesX StagesY 2、StagesZ 1注释明确说明这是共享内存容量限制所致、epilogue 采用TmaWarpSpecialized调度、partial_m 128。最终组合出CollectiveMainloop、CollectiveEpilogue与PersistentTileScheduler产出SsdKernelTmaWarpSpecialized内核。kernel/sm90_ssd_tile_scheduler.hppPersistentTileScheduler实现持久化网格num_blocks B * EH网格大小取min(num_blocks, sm_count)每个线程块通过block_idx gridDim.x步进遍历多个 tile从而在较少 block 下充分利用全部 SM。它同时负责把block_idx解码为(b, eh)与 group 坐标供不同 producer warp 使用。kernel/sm90_ssd_kernel_tma_warpspecialized.hpp内核主体SsdKernelTmaWarpSpecialized其中NumMmaWarpGroups 2硬编码NumLoadWarpGroups 1即每个线程块 3 个 warp groupMaxThreadsPerBlock 3 * NumThreadsPerWarpGroup定义三种WarpGroupRoleProducer、Consumer0、Consumer1Producer 内部再按 lane 划分四种ProducerWarpRoleLoadX、LoadDelta、LoadBC、LoadZ为 X、Delta/DeltaA、B、C、D、Z 以及协作与存储建立了多条 TMA/pipelinePipelineTmaAsync/PipelineTmaStore/PipelineAsync通过warpgroup_reg_alloc/dealloc做寄存器重分配LoadRegisterRequirement 40 - 2*8体现 Hopper 上 producer/consumer 分时复用寄存器的优化手段。device/ssd.hppcutlass::ssd::device::SSDKernel是标准的 CUTLASS 3.x 设备算子门面提供can_implement、get_workspace_size、initialize、update、run与operator()等接口当ArchTag::kMinComputeCapability 90时通过ClusterLauncher发起带 cluster 形状的扩展启动。主程序 111_hopper_ssd.cu 的调用顺序为构造SsdOperation::Arguments含输入输出张量指针与经layout*_transformed()变换的布局→get_workspace_size分配 workspace →can_implement校验 →initialize初始化参数 → warmup默认 3 次→ 正式迭代计时 → 参考实现验证。其中KernelHardwareInfo通过query_device_multiprocessor_count查询 SM 数并传给 tile scheduler。七、主循环四个 BMM 阶段的流水线设计collective/sm90_ssd_gemm_tma_warpspecialized.hpp 中的SsdMainloopTmaWarpSpecialized把 SSD 计算组织为四个矩阵乘阶段每个阶段均复用 CUTLASSCollectiveBuilder生成KernelTmaWarpSpecialized类型的集体算子阶段TileShape布局说明IntraBMM1(L, L, N)NT块内第一阶段矩阵乘对应 B 与 C 张量IntraBMM2(L, D, N)TN块内第二阶段矩阵乘对应累积结果与 XInterBMM1(N, D, L)TN跨块第一阶段对应 B 与 XInterBMM2(L, D, N)NN跨块第二阶段对应状态 P 与 CTileShape (L, D, N) (128, 64, 128)ClusterShape被强制为1x1x1不支持多播。数据通路方面X、B、C 使用 TMA 加载SM90_TMA_LOADAlignment 16/sizeof(Element)保证 16 字节对齐Delta/DeltaA 则使用SM90_BULK_COPY_AUTO批量拷贝。从内核的 Consumer 侧代码可以还原数据流Consumer0Intra 路径等待 B、C →mma_intra_1做 IntraBMM1 →pre_intra_2读取 Delta/DeltaA转换为float后按m n的下三角条件计算expf(deltaA_col - deltaA_row)与 Delta 及 IntraBMM1 结果逐元素相乘即 SEGSUM 的逐元素实现→mma_intra_2用 X 做 IntraBMM2 →store_intra通过协作 pipeline 把部分和交给 epilogue。Consumer1Inter 路径state_init初始化状态张量P含一次经共享内存的转置配合TransposeBarrier/TransformBarrier命名屏障→pre_inter_1等待 B 与 Delta/DeltaA用SM75_U32x4_LDSM_N从共享内存装载 B 片段利用__shfl_sync广播每行最后一个 DeltaA 列值计算expf(last_column - deltaA) * delta * b的加权 B →mma_inter_1用 X 做 InterBMM1 →pre_inter_2把上一轮状态按expf(last_column)衰减后累加进结果并更新状态 →mma_inter_2用状态 P 与 C 做 InterBMM2同时对 DeltaA 取expf供 epilogue 使用 →post_inter_2把新状态经转置写回共享内存。代码注释sm90_ssd_gemm_tma_warpspecialized.hpp特别指出Inter 路径的 DeltaA 处理依赖 GMMA 已知的 TV 分区方式仅对 128x64x128 tile 成立布局为((REG_N,REG_M,REG_N_REP),ATOM_M_REP,ATOM_N_REP)并有static_assert强制GMMA::ALayout_64x16——这正是原文档Only support LxDxN 128x64x128限制的底层原因。八、Epilogue输出张量、状态与前向最终态collective/sm90_ssd_epilogue.hpp 中的SsdEpilogue模板类负责收尾工作维护三条 TMA 流水线EpiloadPipelineD装载对角张量 D、EpiloadPipelineZ装载额外张量 Z、StorePipeline/StorePPipelineTMA 存储 Y 与最终状态 F布尔模板开关HAS_D、D_HAS_HDIM、HAS_Z分别控制是否计算 D、D 是否含 head 维度、是否输出 Z对应共享内存布局与kEpiloadDBytes/kEpiloadZBytes的编译期裁剪输出张量规格在 111_hopper_ssd.cu 中有明确注释x [b, eh, d, c, l]、delta/delta_A [b, eh, c, l]、B/C [b, g, n, c, l]、y [b, eh, d, c, l]、fstate [b, eh, d, n]。epilogue 的store阶段会汇总 Consumer1 产出的tInter2、tDeltaA与更新后的tD按 chunk 写出 Y全部 chunk 完成后由store_p写出最终状态 F。若D_HAS_HDIM为真还需在结尾对 D 流水线做一次consumer_release并推进状态。九、当前限制与性能特征原文档明确列出两项限制结合源码可以给出更精确的边界仅支持 LxDxN 128x64x128TileShape在 111_hopper_ssd.cu 固定为(Options::L, Options::D, Options::N)而 L/D/N 是编译期常量GMMA 布局相关的static_assert进一步把 Inter 路径限定在 128 宽的 tile 上。此外partial_m 128也是写死的。输入输出使用 bfloat16 精度Options::Element cutlass::bfloat16_t累加器为floatDeltaA 张量用float存储而 Delta/B/C/D/Z 均为 bf16。参考实现中的逐元素运算expf等也全部提升到float后计算再回写。性能特征方面原文档给出的三点均有源码佐证利用 TMAX/B/C 走SM90_TMA_LOAD描述符并在启动阶段由单个线程prefetch_tma_descriptors预取描述符warp 特化Producer 按 lane 分工同时推进 X、Delta/DeltaA、B/C、D/Z 的装载两个 Consumer warp group 分别执行 Intra 与 Inter 路径配合PipelineTmaAsync实现多级流水受限于 SEGSUM、ALU boundSEGSUM 涉及共享内存中 Delta/DeltaA 的读写、expf指数运算与逐元素乘加见pre_intra_2、pre_inter_1这些标量浮点操作不经过张量核构成性能瓶颈--use_fast_math编译选项正是为缓解这一 ALU 密集特性。需要注意的是原文档未给出任何具体性能数值本文同样不做任何性能数据推断。十、快速自查清单硬件/工具链确认 GPU 为 SM90compute capability 9.0、CUDA ≥ 12.0、编译器支持 C17构建开启CUTLASS_ENABLE_EXAMPLES并指定CUTLASS_NVCC_ARCHS90a目标111_hopper_ssd运行当前实现仅接受--G1--iterations1会自动进入计时模式并打印 cumsum/ssd 内核耗时与共享内存占用验证默认执行设备端参考实现对比容差 0.05可用--without_verify关闭边界tile 固定为 128x64x128精度为 bf16 输入 / fp32 累加。版权说明本示例源码与文档版权归 NVIDIA CORPORATION AFFILIATES采用 BSD-3-Clause 许可证SPDX-License-Identifier: BSD-3-Clause具体条款见 README.md 与各源文件头部使用与分发时请遵守该许可证要求。【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考