Warp Specialization在 GEMM Kernel 中的性能建模与优化：仿真工具与 ≤21.5% 最大预测误差

关键词：Warp 特化、GPU 性能建模、GeMM-WS 内核、微分方程、SMT 求解器

• A Performance Model for Warp Specialization Kernels
• https://arxiv.org/pdf/2506.11209
• 1 万字，阅读需 33 分钟，播客 18 分钟

• ISCA'25 建模GPU LLM推理性能AMALI: 借助张量核心优化与warp分布，推理误差低于GCoM的GPU分析模型
• LLMPerf：GPU性能建模与大型语言模型的融合
• Fake Runs, Real Fixes！复用GRM的AI加速器微架构分析方法：15%All-Gather和4%LLM性能提升

本文提出了一种针对 warp 特化内核的性能模型，聚焦 warp 大小、分块尺寸、输入矩阵规模、内存带宽及线程分歧等关键因素，通过微分方程建模并经仿真与实验验证，可精准预测内核执行时间。

模型以 warp 为调度最小单元，将 DMA warp 视为生产者、MATH warp 视为消费者，从 wave 和 stage 层面分析 GeMM-WS 内核执行流程，构建了描述 DMA 加载与 MATH 计算事件时序的微分方程，还考虑了共享内存循环缓冲区的约束。

在 NVIDIA A6000 GPU 上的测试显示，求解器多数场景下数分钟内可获最优解，模型对 2048 种配置的执行时间预测平均误差仅 4.5%。该模型深化了对 warp 特化技术的理解，也为 GPU 加速应用的编译器优化、内核参数调优和算法设计提供了实用依据。

交流加群请在 NeuralTalk 公众号后台回复：加群

unsetunset本文目录unsetunset

• 本文目录
• 关键问题
• 问题 1. Warp 特化性能模型跨架构迁移的精度保持与参数重构问题
• 问题 2. 多目标优化下 SMT 求解器的解空间爆炸与模型适配调整问题
• 一、背景
• 1.1 通用矩阵乘法
• 1.2 线程束专用化（Warp Specialization）
• 1.3 CUTLASS 库中的线程束专用化 GeMM 核
• 二、性能模型
• 2.1 高层性能分析
• 2.2 一个波次内的事件分析
• 三、寻找最优配置
• 3.1 执行模拟
• 3.2 利用 SMT 求解器寻找最优解
• 3.3 模拟工具与求解器的交叉验证
• 四、评估
• 4.1 求解时间
• 4.2 预测精度
• 五、结论
• 参考文献

交流加群请在 NeuralTalk 公众号后台回复：加群

unsetunset关键问题unsetunset

问题 1. Warp 特化性能模型跨架构迁移的精度保持与参数重构问题

该 Warp 特化内核性能模型基于 NVIDIA A6000 GPU 验证，若迁移至架构差异显著的 GPU（如 AMD MI300 或 NVIDIA H200），其核心的微分方程建模逻辑及 4.5%的预测精度能否保持？模型中哪些参数需针对性重构？

该 Warp 特化内核性能模型迁移至架构差异显著的 GPU 时，核心微分方程的建模逻辑框架可保留，但 4.5%的预测精度无法直接保持，且多个硬件相关参数需针对性重构。

1. 建模逻辑的保留性：模型的核心是基于生产者-消费者（DMA/MATH warp）的事件时序微分方程、wave/stage 的执行拆解逻辑，以及共享内存循环缓冲区的约束规则，这部分是针对 Warp 特化内核的通用调度规律设计，不与 NVIDIA A6000 架构强绑定，因此迁移后核心逻辑仍适用。
2. 预测精度的变化：4.5%的平均预测误差是基于 NVIDIA A6000 的实测参数（如 DMA 吞吐量 152.96 MB/s、启动延迟 0.770 us，ALU 吞吐量 24.61 GB/s）得出的，不同 GPU 的硬件吞吐量、启动延迟、SM 数量等存在显著差异，直接复用参数会导致预测偏差，精度无法维持。
3. 需重构的参数：
• 计算类参数：ComputeThrough（计算吞吐量）、LoadThrough（加载吞吐量）、ComputeStartupLatency/LoadStartupLatency（计算/加载启动延迟）需重新实测；
• 架构类参数：SM 数量会影响 wave 数计算公式  ，共享内存容量会改变循环缓冲区最大 tile 数   的取值，进而影响微分方程中   和   的约束；
• 调度类参数：若 GPU 的 Warp/Wavefront 大小不同（如 AMD MI300 为 64，NVIDIA H200 为 32），需调整 Warp 特化的分工粒度，同步机制（如信号量开销）的差异也需纳入模型参数。

问题 2. 多目标优化下 SMT 求解器的解空间爆炸与模型适配调整问题

论文以“最小化执行时间+MATH Warp 等待时间”为优化目标，未提及功耗、内存占用等实际部署约束，若引入多目标优化场景，基于 Z3 SMT 求解器的优化器是否会出现解空间爆炸？模型需如何调整以平衡求解效率与解的实用性？

引入功耗、内存占用等实际部署约束的多目标优化后，基于 Z3 SMT 求解器的优化器【存在】解空间爆炸的风险，需从约束设计、求解策略、编码优化三方面调整模型以平衡效率与实用性：

1. 解空间爆炸的必然性：论文中优化器仅以“最小化执行时间”和“最小化 MATH Warp 等待时间”为目标，编码采用 Array 理论与整数类型简化计算，且测试的分块配置仅为 64/128、矩阵规模 ≤1024。引入功耗、内存占用约束后，需新增“功耗-分块尺寸”“内存占用-缓冲区大小”等关联方程，变量维度与约束条件大幅增加，若矩阵规模或分块选项扩展，易出现解空间爆炸，导致求解时间从“分钟级”急剧延长。
2. 模型调整方向：
• 约束优先级重构：将功耗、内存占用设为“硬约束”（如共享内存占用不超过硬件上限），而非独立优化目标，仅保留执行时间为核心优化目标，过滤无效解空间；
• 解空间预剪枝：利用论文中的仿真工具，先对分块配置进行快速筛选（如排除内存占用超标的   组合），再将候选配置输入 SMT 求解器，减少求解器的计算量；
• SMT 编码优化：在原有 Array 理论编码基础上，引入启发式规则缩小变量取值范围（如分块尺寸仅取硬件适配的 2 的幂次），并采用 Z3 的加权多目标优化功能，将功耗、内存占用转化为惩罚项融入执行时间目标，避免多目标的解空间膨胀；
• 分层求解策略：先通过仿真工具得到近似最优解的范围，再让 SMT 求解器在该范围内精细化搜索，兼顾求解效率与解的实用性。

unsetunset一、背景unsetunset

1.1 通用矩阵乘法

通用矩阵乘法核（General Matrix Multiply kernel，简称 GeMM）是一种用于执行矩阵乘法的专用计算例程，常用于数值计算以及机器学习、科学计算等各类应用中。

GeMM 核接收两个输入矩阵 A 和 B，并生成一个输出矩阵 C，满足：

其中，矩阵 A 的维度为  ，矩阵 B 的维度为  ，输出矩阵 C 的维度为  。

在 GeMM 的场景下，分块（tiling） 是一种通过将输入矩阵划分为更小的、尺寸规则的子矩阵，并基于这些子矩阵执行乘法（而非直接操作原始大矩阵）来提高数据局部性（指数据在内存中存储的位置越集中，CPU/GPU 访问数据的速度越快）的技术。

图 1 展示了分块的一个示例。在本文的研究中，我们使用一个由三个非零整数组成的元组来表示分块配置，记为  。在该分块配置   中，第一个输入矩阵 A 被划分为维度为   的分块，第二个输入矩阵 B 被划分为维度为   的分块，输出矩阵 C 则被划分为维度为   的分块。

1.2 线程束专用化（Warp Specialization）

GeMM 核的实现通常采用串行代码逻辑：先从全局内存（GPU 中容量较大、访问速度较慢的内存，供所有线程共享）加载数据，接着执行矩阵乘法计算，最后将结果写回全局内存。这种实现方式存在一个缺陷：计算操作与数据访问操作是同步执行的，这会导致算术逻辑单元或直接内存访问单元出现空闲时间。

线程束专用化核（Warp Specialization Kernels） [1] 是一类基于每个线程的 ID 动态分支，并让每个线程束专门执行某一特定任务的核。

在 GeMM 场景下，每个线程束被专用化以执行矩阵乘法的不同阶段。通过线程束级别的专用化，若调度得当，数据访问与计算操作可并行执行，同时避免因线程束内分支导致的性能下降。

1.3 CUTLASS 库中的线程束专用化 GeMM 核

本文的研究聚焦于 NVIDIA CUTLASS 库[2]中实现的线程束专用化 GeMM 核（记为 GeMM-WS）。

GeMM-WS 核从全局内存加载输入矩阵、执行矩阵乘法，并将结果写回全局内存。GeMM-WS 核的伪代码如算法 1 所示。

在 GeMM-WS 中，每个线程束通过其功能实现专用化，存在两类专用化线程束：

我们将 DMA 线程束加载数据的过程称为前序阶段（prologue），将计算线程束写回结果的过程称为收尾阶段（epilogue）。每个计算线程束（MATH warp）关联 1 个或 2 个 DMA 线程束：

• 在“1 个计算线程束关联 1 个 DMA 线程束”的配置下：由于 GeMM 是二元运算（需两个输入矩阵 A 和 B），该 DMA 线程束会加载两个输入分块（分别来自 A 和 B）；
• 在“1 个计算线程束关联 2 个 DMA 线程束”的配置下：每个 DMA 线程束仅加载一个输入分块（一个加载 A 的分块，一个加载 B 的分块）。

GeMM-WS 强制实施两类同步机制：

1. 核中实现了一种“等待-信号”信号量（wait-signal semaphore，一种同步机制：线程需先“等待”（wait）信号量就绪，再执行操作，操作完成后发送“信号”（signal）通知其他线程），用于确保 DMA 线程束与计算线程束之间的“先发生后执行”（happen-before）关系（即 DMA 加载数据的操作必须在计算线程束使用该数据之前完成）；
2. 任何时刻最多只能有一个计算线程束处于活跃状态——若存在多个活跃的计算线程束，可能导致负载不均衡（部分 SM 忙碌、部分 SM 空闲），进而降低整体性能。

unsetunset二、性能模型unsetunset

我们构建了一个用于预测线程束专用化核执行行为的性能模型。该模型中最小的调度单元是线程束（warp），并将 DMA 线程束称为“生产者（producer）”，计算线程束称为“消费者（consumer）”。模型假设：任何时刻最多只能有一个计算线程束或一个 DMA 线程束处于活跃状态。

2.1 高层性能分析

本节从高层视角分析 GeMM-WS 核的性能。

2.1.1 波次（Wave）

在 GeMM-WS 核中，每个输出分块的计算被映射到一个协同线程数组（Cooperative Thread Array，简称 CTA，即 GPU 中的“线程块”，由多个线程组成，可在 SM 上执行）。核启动时分配的 CTA 数量受两个因素限制，取二者中的较小值：

• 所需的输出分块总数；
• 系统中 SM 的数量。

• 若输出分块数量小于或等于系统中的 SM 数量，我们可将每个输出分块的计算分配给一个独立的 SM——此时每个 SM 最多只需处理一个分块，计算可在一个波次（wave） 内完成（波次指 SM 执行 CTA 的批次，一批次内所有 SM 同时处理一个 CTA）。
• 若输出分块数量大于 SM 数量，则计算无法在一个波次内完成。例如：当有 256 个输出分块需处理，且系统有 84 个 SM 时，分块会被均匀分配给各 SM——其中 80 个 SM 各分配 3 个分块，剩余 4 个 SM 各分配 4 个分块，整体需 4 个波次才能完成计算。

一般而言，完成 GeMM 计算所需的波次数量可表示为：

输出分块总数可通过 M 维度和 N 维度上的分块数量相乘得到：

GeMM-WS 核的总执行时间等于所有波次的执行时间之和，再加上第一个波次启动前的初始化时间（ ）：