一、量化分析基础:优化前先算账
1. Amdahl 定律:加速比的天花板
Amdahl 定律量化并行计算中加速比与串行比例的关系:
\[ \text{Speedup} = \frac{1}{1 - P + \dfrac{P}{N}} \]式中 \(P\) 为并行部分比例,\(N\) 为并行部分的加速比。示例:若某程序 80% 的代码可并行运行,对于这部分的加速比无穷大(\(N \to \infty\)),则此程序加速比的上界为 \(1/(1-0.8)=5\times\)。
这个公式反映的是,堆砌 CPU 核或 GPU ALU 对应用加速比的提升也不是无限的。在实际 E2E(End-to-End)性能优化过程中,若只对某部分进行加速会出现边际效用递减,因此优化过程中需要不断重新采样,始终抓住主要矛盾(即当前的性能瓶颈点)。
2. Roofline 模型:判断计算受限还是访存受限
Roofline 模型通过直观的图形化方式,将程序的计算强度与系统的内存带宽、计算能力相结合,帮助快速定位程序性能的限制因素。横轴是计算强度(Operational Intensity,单位 FLOP/Byte),纵轴是性能(通常以 FLOPS 衡量)。模型由两条关键曲线构成:计算能力曲线(水平线,代表硬件理论峰值 \(\pi\))和内存带宽曲线(斜线,斜率为内存带宽 \(\beta\)),二者组合形成类似屋顶的形状。
▲ Attainable Performance (FLOP/s)
│ ______________ ← 计算能力曲线(水平线,π = 峰值算力)
│ / Compute Bound
│ /
│ Memory /
│ Bound /
│ /
│ β (斜率=内存带宽)
│ /
└─────────────┴──────────────────► Operational Intensity (FLOP/Byte)
I_max = π / β
计算能力曲线由处理器的时钟频率、每个时钟周期可执行的浮点运算数量等硬件参数决定。例如一款处理器每周期执行 8 次浮点运算,时钟频率 3GHz,理论峰值计算能力就是 24 GFLOPS。内存带宽曲线反映系统数据传输速率上限,在内存受限区(Memory-Bound),理论峰值性能遵循:
\[ P(\text{性能}) = \beta(\text{内存带宽}) \times I(\text{计算强度}) \]这意味着在该区域内,程序每秒完成的计算量直接受限于访存速度;随着计算强度提升,可达到的性能上限沿斜线线性增长,直至触及硬件算力天花板 \(\pi\)。
应用场景:① 算法优化——判断算法受限于计算能力还是内存带宽,例如矩阵乘法若性能点位于内存带宽曲线附近,说明需要通过优化数据布局、增加数据缓存来提升性能;② 硬件评估——对比不同硬件平台对同一程序的性能表现,不同处理器/内存组合会产生不同形状的"屋顶"。局限:基于理想硬件模型和简化假设,忽略了缓存层次结构(可针对存储层次与不同算力单元分别建模来修正)、硬件流水线与访存延迟等影响。
二、计算优化:让计算单元持续高效运转
1. 调控 Occupancy
定义:GPU 上的 Occupancy 指一个 MP 的 Active Warp 数量与该 MP 最多支持的线程束数量之比:
\[ \text{Occupancy} = \frac{\text{每个MP上的Active Warp数量}}{\text{MP最多支持的Warp数量}} \]Active Warp 指 MP 上同时被调度执行的线程束。一个 MP 上的最大线程束数量固定,由 Warp scheduler 决定,实际 Active Warp 数量由 Kernel 中硬件资源使用量所限制。
算例(教材原例):每个 Block 有 32768 个寄存器和 32 KB 共享内存,每 Block 支持最多 8 个 warp;MP 总资源为 131072 个寄存器和 192 KB 共享内存。按寄存器计算最多可驻留 4 个 Block(131072 / 32768 = 4);按共享内存计算则可驻留 6 个 Block(192KB / 32KB = 6);每个 MP 上实际驻留 Block 数取最小值 \(min(4,6)=4\),因此 Active Warp = 4×8 = 32;Occupancy = 32/96 = 1/3(96 为该 MP 最多支持的线程束数量)。公式可概括为:Occupancy = min(按共享内存可驻留的Active Warp数, 按寄存器可驻留的Active Warp数) / MP最大Warp数。
当一个线程束因等待数据而停止执行时,MP 可以调度另一个线程束来执行,以降低延迟。因此,Active Warp 数量的增加可以提高 MP 的利用率,从而提高指令吞吐量。高 Occupancy 的好处:① 确保总有足够的 Active Warp 用于切换,隐藏长延迟操作;② 足够大的全局内存访问并发数,提高内存带宽利用率;③ 确保每个 MP 有足够多线程平衡负载,提高负载均衡性能。
2. 减少线程束分歧(Warp Divergence)
在 MT GPU 中,一个线程束共享同一个程序计数器(PC),但 SIMT 编程模型下一个线程束内不同线程又允许执行完全不同的任务,这种情况称为线程束分歧(Warp Divergence)。条件控制的分支语句(如 if-else、for-loop、do-while)可能引发 Warp Divergence:
if (tid < 31) {
// task A
} else {
// task B
}
假设只有一个线程束共 32 个线程,前 31 个线程进入 if 分支执行 task A,最后一个线程进入 else 分支执行 task B,线程束内的线程在这里分叉。由于架构上线程束只有统一的 PC,GPU 会分别串行地执行 if 和 else 分支,并在执行时为每个线程设置相应的活跃掩码,确保只有活跃的线程才真正执行。当发生 Warp Divergence 时,线程束发射的指令没有充分利用 ALU 的吞吐量,整个线程束完成任务的时间近似变成各分支任务时间之和。
规避策略:① 重新设计算法,让一个线程束内的任务统一;② 尝试使用三元表达式等逻辑,将条件判断与赋值操作融合,避免显式分支。
案例:并行规约中的 Divergence 优化
假设 GPU 的 warp size 为 2,启用 4 个线程束计算,数据在共享存储器中已就绪。版本一(含 Warp Divergence):
for (int stride = 1; stride < 8; stride *= 2) {
if (tid % (2 * stride) == 0) {
smem[tid] += smem[tid + stride];
}
__syncthreads();
}
当 stride=1 时,每个线程束都有 Divergence 且都只有一个线程激活;stride=2 和 stride=4 时同理。版本二(改变算法流程以减少 Divergence):
for (int stride = 4; stride > 0; stride /= 2) {
if (tid < stride) {
smem[tid] += smem[tid + stride];
}
__syncthreads();
}
虽然最后一轮的线程束内累加无法避免 Divergence,但 stride=4 和 stride=2 时,每个线程束的执行都是统一的(要么整个 warp 全活跃,要么全不活跃)。版本一按"取模"分支,激活线程在每个 warp 内部呈"稀疏交错"分布;版本二按"tid < stride"分支,激活线程集中在低编号的连续区间内,因此更容易让整个 warp 要么全激活、要么全不激活,显著减少每个 warp 内部的分歧开销。这是并行规约优化的经典范式,也是后文 Reduce 案例的基础。
3. 指令级并行(ILP)
指令级并行是一种常见的性能优化手段,其中流水线技术是实现它的重要方法。几乎所有处理器都通过流水线让指令重叠执行以提高性能,这种指令之间的并行执行称为指令级并行(Instruction-Level Parallelism, ILP)。
GPU 流水线完成取指、译码后,会进行相关的依赖检查(检查结构冒险、数据冒险等是否解除),然后将指令发射到不同硬件功能单元执行。要最大限度减少流水线停顿,开发者需要在编写程序时通过指令调度来规避硬件冲突。其中数据冒险,尤其是 RAW(Read-After-Write)这样的真依赖,涉及长延迟访存指令时需格外小心:从单个线程束的 ILP 视角看,需要精心排布软流水——在消费写回数据之前,尽量穿插其他无关指令,让不同硬件功能单元可以并行执行,避免流水线停顿。
Load Gmem ──────────────────────────────► Store Smem
Load Smem Load Smem Load Smem Load Smem Load Smem Load Smem
↓ ↓ ↓ ↓ ↓ ↓
ALU ALU ALU ALU ALU ALU ──► sync
以经典的 SIMT 单精度矩阵乘法为例,通过精心排布的流水线,计算和访存操作可以并行执行,避免不必要的数据依赖,使 ALU 的计算不被打断。教学要点:ILP 优化的核心思路是——"当一条长延迟指令(如访存)发出后,不要让流水线空等结果,而是立即安排其他不依赖该结果的指令去填充这段等待时间",这与 CPU 的乱序执行/软件流水线思想一致,只是在 GPU 上更多依靠编译器调度和程序员显式排布指令顺序来实现。
4. Warp Specialization:让线程束各司其职
定义:Warp Specialization 是一种流水线并行的优化技巧,线程束被划分成不同角色,例如负责数据搬运的"访存线程束"、负责逻辑运算的"计算线程束"、负责任务分发的"调度线程束",这些角色之间形成紧密的生产者-消费者关系。
优势:① 不同功能角色天然独立调度,每个角色只需关注自己的生产-消费关系(数据冒险),硬件功能单元可以在冒险解除后尽快运行,不必担心被不相干指令阻塞;② 单个角色内寄存器压力降低,更有助于编译器进行寄存器分配,不同线程束的独立性带来天然的指令级并行,把编译器指令调度排布的压力 offload 到硬件调度器上。
数据搬运优化:角色划分后,如果沿用原本的 LSU 数据搬运方式只有一部分线程参与,同一线程束需要发射的指令数量会大大增加,也可能带来更多寄存器开销。从寄存器角度考虑,硬件架构最好能直接将数据从全局内存搬运到共享存储器,避免寄存器问题引起的不必要 RAW 依赖;MT GPU 支持 TME 功能,可以用一个线程束完成大量数据的搬运,以节省资源。
同步机制:异步事务屏障(Async Transaction Barrier)。角色划分后,原本独立的线程束之间需建立高效的生产-消费者通信。传统的 Shared Memory 或常规 Barrier 同步机制耦合度过高,难以应对复杂的异步流水线。为此,MUSA GPU 引入了异步事务屏障能力:该机制实现了同步粒度的精细化控制,还支持与TME(张量管理引擎)协同,在数据异步搬运的同时解耦线程同步,显著提升硬件利用率。
生产者 (Producer) 消费者 (Consumer)
┌──────────────┐ ┌──────────────┐
┌───►│producer acquire│◄──────┐ ┌─────►│consumer wait │
│ └──────┬───────┘ │ │ └──────┬───────┘
│ ▼ │ │ ▼
│ ┌──────────────┐ │ │ ┌──────────────┐
│ │ TME Load │ │ │ │ ALU │
│ └──────┬───────┘ │ │ └──────┬───────┘
│ ▼ │ │ ▼
│ ┌──────────────┐ │ │ ┌──────────────┐
└────┤producer release├──────┘ └──────┤consumer release│
└──────────────┘ └──────────────┘
核心循环中,通过两个一组的 Async Transaction Barrier 相互通知生产者和消费者的关系:① 生产者尝试通过 producer acquire 获取共享存储器的写入权限,该内存是否允许写入依赖消费者在完成数据消费后通过 consumer release 通知生产者;② 允许写入后利用 TME 进行数据搬运,由于 TME 是异步指令,生产者无需等待搬运完成,只需在发射相关指令后通过 producer release 明确已完成指令发射(对应 Async Transaction Barrier 的 arrive,此时数据不一定已读入);③ 消费者等待生产者数据返回,TME 每完成一部分数据搬运就在 Async Transaction Barrier 里记录,直到已完成的搬运量达到设定量,consumer wait 屏障才允许通过读取数据计算,完成后再通知生产者数据已释放。
if (threadIdx.x % 2 == 0) 分支的核函数代码,判断是否存在 Warp Divergence,并给出改写方案。③ 概念辨析:为什么"高 Occupancy"不等于"高性能"?请结合 Compute Bound 场景举例说明。三、内存优化:把带宽榨干
1. 合并访存(Coalesced Access)
在 GPU 编程中,合并访存指线程束中的线程按特定规则访问全局内存时,硬件能够将这些分散的内存访问请求合并为内存事务(Transaction),从而显著提高内存访问效率的机制。
MTT S5000 的合并粒度:在 MTT S5000 的一个线程束内,会以每 4 个或 8 个线程为粒度,经过 Coalescing Unit,将这些线程的并发访问合并为多个 128 字节(L1/L2 的 cacheline size)的请求,之后内存子系统的下游都以 128 字节的事务为单位进行操作。
反例(低效访存):为了尽量少产生额外事务,连续线程访问的数据块应尽量连续且对齐。例如对于 float* src,如果使用 src[threadIdx.x * 32] 进行读取,带宽利用率仅有 4 字节/128 字节 = 1/32,十分低效。
2. 向量化访存(Vectorized Load/Store)
向量化访存即使用一条指令同时请求多个数据。CUDA 最高支持 128 比特的向量化访存,即一条指令请求 4 个 float 数据、8 个 half 数据或 16 个 int8 数据;MTT S5000 GPU 支持更大的 burst 数据量请求,单条指令最高可以进行 1024 比特的数据访问。
为什么合并访存之外还需要向量化?在满足合并访存的情况下,固定的访问数据量会产生固定数量的 128 字节事务,但向量化访存产生的访存指令数量要少于标量访问,而处理访存指令的单元(LSU)的处理能力有上限,若发射的访存指令数量超出 LSU 处理上限,会导致访问请求无法下发,影响带宽利用率。若单个访存请求小于 4 字节,会使访存效率大幅降低,此时建议使用 Vectorize Load/Store 优化。
// 标量访问(4 条指令,每条 4 字节)
float a0 = src[idx];
float a1 = src[idx+1];
float a2 = src[idx+2];
float a3 = src[idx+3];
// 向量化访问(1 条指令,16 字节,float4)
float4 a = *((float4*)(src + idx));
3. 减少共享存储器的存储体冲突(Bank Conflict)
定义:线程间共享存储器的存储体冲突(Bank Conflict)指线程束中多个线程同时访问共享存储器中的同一存储体,导致内存访问串行执行、降低访问效率的现象。共享存储器的存储体设计是为了实现并行访问,而存储体冲突会破坏这种并行性,成为性能瓶颈之一。
| 冲突程度 | 描述 | 效率 |
|---|---|---|
| 无冲突 | 32 个线程访问 32 个不同的存储体,并行处理 | 访问效率最高 |
| 部分冲突 | 多个线程访问同一存储体(如每两个线程访问同一存储体) | 访问延迟些许增加,但所提供的带宽会大大减少 |
| 全冲突 | 线程束中所有 32 个线程访问同一存储体,需进行 32 次访问 | 效率最低 |
广播机制:现代 GPU 中同一线程束访问共享存储器中同一地址时,通常会有广播机制,只需进行一次访问——这与"访问同一存储体的不同地址"是不同的情况,需明确区分:同地址广播不算冲突,同存储体不同地址才算冲突。
两种解决方案:① 增加 Padding——使每次访问共享存储器请求错开一个存储体,但会浪费一些容量;② 利用地址散列或 Swizzle 地址映射——在工程实践中更推荐方案二。
# 根据访问的数据调整函数中的 mask 和 shift 参数
def fast_swizzle(self, lms_addr):
temp = (np.uint32(lms_addr) & self.mask) >> self.shift
return np.uint32(temp) ^ np.uint32(lms_addr)
self.mask、self.shift 的具体取值依赖硬件存储体数量与访问粒度,教材正文未展开推导过程,只给出函数框架。本课不猜测补全具体数值,建议结合原书 32×32 索引矩阵示意图理解其"通过异或重排地址、打散规律性冲突模式"的核心原理,实际取值请按目标硬件规格计算。data[threadIdx.x] 与 data[threadIdx.x * 4] 哪种访问模式更符合合并访存?为什么?② 计算题:若一个线程束访问共享内存时每 2 个线程映射到同一存储体(部分冲突),相比无冲突情形,理论上访问所需周期数会增加多少倍?③ 设计题:给出一个会产生 Bank Conflict 的共享内存访问模式(如按列访问二维数组的转置场景),分别用 Padding 和 Swizzle 给出改写思路。四、多 GPU 编程:单机多卡与多机多卡
1. 单机多卡(Multi-GPU)
单机多卡的核心逻辑是"分而治之",即将任务拆分为若干子任务分配给不同 GPU 独立完成,再通过跨卡通信完成结果汇总。
(1) 设备枚举与选择:
int device_count;
musaGetDeviceCount(&device_count);
musaSetDevice(device_id); // 指定当前线程操作的设备
可结合多线程或多进程方式,让每个执行单元独立管理一张卡。
(2) 数据分布与管理:将输入数据按任务切分,通过 musaMemcpy() 分别复制到各设备的全局内存,可根据显存容量与计算量动态调整分片大小,确保负载均衡。
(3) 计算与通信:在每张卡上独立启动 Kernel 执行分配的子任务;跨卡通信借助 MUSA P2P 功能,通过 musaMemcpyPeer() 实现设备间数据复制以降低主机参与的开销,使用前需通过 musaDeviceEnablePeerAccess() 启用 P2P 功能:
// P2P 通信基本模式(依据原书文字描述重构的安全示例)
musaSetDevice(0);
musaDeviceEnablePeerAccess(1, 0);
musaSetDevice(1);
musaDeviceEnablePeerAccess(0, 0);
// 之后可直接:
musaMemcpyPeer(dst_ptr, dst_device, src_ptr, src_device, size);
musaMemcpyPeer() 与 musaDeviceEnablePeerAccess() 的作用及调用先后关系,未给出完整代码清单,以上代码为根据教材文字重构的安全示例,实际参数顺序请以当前 MUSA 官方 API 文档为准。性能优化策略:① 负载均衡——先量化各 GPU 的算力/带宽上限,再根据子任务"计算强度"动态分配工作量,避免部分 GPU 过载、部分空闲;② 消除通信瓶颈——硬件层面用 MTLink 或 PCIe Gen4/5 提高跨卡带宽,软件层面将数据通信与计算重叠执行,利用 MUSA 流(musaStream_t)在异步复制与核函数执行间交替调度,实现"边传输、边计算";③ 内存优化——使用统一内存简化数据管理但需谨慎分析 Page Migration 延迟,使用页锁定内存在主机与设备间复制以提高带宽;④ 使用高效通信库——MCCL 提供针对多 GPU 间 all-reduce、broadcast 等常见通信模式的高度优化实现,无须手动编写复杂跨卡通信逻辑。
工业案例:GROMACS(分子动力学仿真)推荐用多线程方式做单机多卡,编译时设置 -DGMX_THREAD_MPI=ON,运行时若不设置 gpu_id 则 GPU 数量默认为 ntmpi 参数指定的数量;设置环境变量 GMX_ENABLE_DIRECT_GPU_COMM=1 可让 MT GPU 间通过 MTLink 进行 P2P 通信。Amber(生物分子动力学)通过 MPI 管理多卡协同,MUSA 版本支持多卡 pmemd.musa.MPI,内部通过 MCCL 通信,单节点内最多支持 8 卡;由于单次 MD 跨卡扩展收益有限,通常推荐在单节点多 GPU 上同时执行多个独立任务以获得更高总吞吐量。
2. 多机多卡(集群级多 GPU)
多机多卡的核心逻辑是"分布式协同",将任务拆分至不同节点的 GPU,通过跨节点通信完成数据同步,底层依赖 MUSA 设备管理逻辑,但需额外考虑节点间互联、分布式数据一致性等问题。
GROMACS 多机多卡:编译时指定 MPI 相关参数即可,-ntmpi 指定运行节点数,-ntomp 指定每个节点的线程数,-nt 指定总线程数。
工业案例:电化学硬件求解器多机多卡。某电池厂商基于 MUSA 生态开发的电化学硬件求解器验证了多机多卡加速能力:整个运行程序使用 Docker 镜像分发,GPU Kernel 部分用 C++ API 包装后由 Python 脚本调用。核心步骤:① 构建包含 MUSA 驱动、SSH、MPI 的 Docker 镜像,配置 SSH 密钥对与统一监听端口,确保各节点环境一致;② 在多节点测试 MPI 样例程序,确保可以进行 MPI 互联;③ 在镜像中安装 Python 与 MPI4Py;④ 通过 MPI4Py 实现节点间数据通信,完成业务代码联调。验证结果:该方案在12 台服务器(每台 8 个 MUSA GPU,共 96 卡)的集群中测试验证,跨节点通信无异常,算力随 GPU 数量线性增长。
musaDeviceEnablePeerAccess()?如果不调用会发生什么?② 设计题:给定 4 张算力不均等的 GPU(算力比例 1:1:2:2),设计一个任务切分策略使总体吞吐量最大化。③ 讨论题:多机多卡场景下,为什么"先验证 MPI 样例程序互联,再接入业务代码"是推荐的工程实践顺序?五、Step-by-Step 实验:四大性能优化案例
本节精选 4 个典型场景(Reduce、GEMV、Flash Attention、GEMM),覆盖"简单规约 → 访存受限 → 长序列 → 计算密集",将优化策略融入实践。
案例 1:Reduce 优化
问题定义:Reduce 计算 \(a = a_0 \oplus a_1 \oplus a_2 \oplus \cdots \oplus a_n\),其中 \(\oplus\) 可以是任意可规约运算符(加法、求最大值等),本节采用加法运算符。
分治策略:① 将输入数据分割成多个块,每个块分配给一个线程块;② 在每个线程块内执行局部 Reduce,将单个数据块缩减到一个值;③ 将所有局部结果合并为最终结果。
Level 0: 8 blocks ▽▽▽▽▽▽▽▽ (每个 block 内部先做局部规约)
\ \ \...
Level 1: 1 block ▽ (再对 8 个局部结果做一次规约)
Step 1:基础实现(共享内存版),可直接编译运行:
template<int THREAD_PER_BLOCK>
__global__ void ReduceLarge(float *in, float *out_partial, int num_ele) {
__shared__ float sdata[THREAD_PER_BLOCK];
unsigned int tid = threadIdx.x;
unsigned int bid = blockIdx.x;
// 从全局内存加载数据到共享内存
unsigned int offset = bid * THREAD_PER_BLOCK + tid;
sdata[tid] = in[offset];
__syncthreads();
// 在共享内存中执行归约
for (int s = THREAD_PER_BLOCK / 2; s >= 2; s >>= 1) {
if (tid < s) {
sdata[tid] += sdata[tid + s];
}
__syncthreads();
}
// 写出结果
if (tid == 0) {
out_partial[bid] = sdata[0] + sdata[1];
}
}
// 启动内核(示例参数)
int num_ele = 40960;
constexpr int THREAD_PER_BLOCK = 1024;
int block_dim = THREAD_PER_BLOCK;
int grid_dim = num_ele / block_dim;
ReduceLarge<THREAD_PER_BLOCK><<<grid_dim, block_dim>>>(in_ptr, out_partial_ptr, num_ele);
注意循环条件是 s >= 2(而非 s >= 1),最后一步 sdata[0] + sdata[1] 在写出结果时单独完成,本质上等价于把最后一次规约展开处理。共享存储器是一个 block 内线程共享的内存,访存带宽远高于全局内存、延迟远低于全局内存,所以先将全局内存数据加载到共享存储器再执行 Reduce,可避免对全局内存的频繁访问。
Step 2:Shuffle 指令优化。共享存储器相较全局内存有较低访问延迟,但相较寄存器仍高一个数量级。Shuffle 指令最重要的特性是线程束内的寄存器可以相互访问并使用,因此在支持 shuffle 指令的硬件上,可先用 shuffle 指令完成线程束内的 Reduce,再借助共享存储器完成各线程束中间结果的 Reduce,从而大量减少对共享存储器的访问:
template <int blockSize>
__device__ __forceinline__ void WarpReduce(float& rv1) {
if constexpr (blockSize >= 32) {
rv1 += __shfl_down_sync(0xffffffff, rv1, 16);
}
if constexpr (blockSize >= 16) {
rv1 += __shfl_down_sync(0xffffffff, rv1, 8);
}
if constexpr (blockSize >= 8) {
rv1 += __shfl_down_sync(0xffffffff, rv1, 4);
}
if constexpr (blockSize >= 4) {
rv1 += __shfl_down_sync(0xffffffff, rv1, 2);
}
if constexpr (blockSize >= 2) {
rv1 += __shfl_down_sync(0xffffffff, rv1, 1);
}
}
Step 3:工作负载微调。① vector load(向量加载)——增加单个线程处理的任务量,减少启动的 block 数目,对访存更友好也减少指令条数;② fine tune block size——根据数据规模与不同硬件调优线程块大小以获取更好性能。
Step 4:小规模场景的单 Kernel 方案。分治策略需要启动两个 Kernel,适用于大规模情况;小规模情况下想取得较好性能,可仅用一个线程块计算,一步得出最终结果,降低对全局数据访问的开销:
template<int THREAD_PER_BLOCK, int Vlen>
__global__ void ReduceSmall(float *in, float *out, int num_ele) {
using vec = float4;
float src_reg[Vlen];
__shared__ float sdata[THREAD_PER_BLOCK];
// 每个线程从全局内存加载 Vlen 个元素
unsigned int tid = threadIdx.x;
unsigned int offset = tid * Vlen;
*((vec*)src_reg) = *(vec*)(in + offset);
float sum = src_reg[0];
#pragma unroll
for (int i = 1; i < Vlen; ++i) {
sum += src_reg[i];
}
// 循环加载剩余数据
for (int index = offset + THREAD_PER_BLOCK * Vlen; index < num_ele;
index += THREAD_PER_BLOCK * Vlen) {
*((vec*)src_reg) = *(vec*)(in + index);
#pragma unroll
for (int i = 0; i < Vlen; ++i) {
sum += src_reg[i];
}
}
sdata[tid] = sum;
__syncthreads();
// 在共享内存中执行归约
for (int s = THREAD_PER_BLOCK / 2; s >= 2; s >>= 1) {
if (tid < s) {
sdata[tid] += sdata[tid + s];
}
__syncthreads();
}
// 写出最终结果
if (tid == 0) {
sum = sdata[0] + sdata[1];
out[0] = sum;
}
}
// 启动内核
int num_ele = 4096;
constexpr int THREAD_PER_BLOCK = 512;
constexpr int vlen = 4;
int block_dim = THREAD_PER_BLOCK;
int grid_dim = 1;
ReduceSmall<THREAD_PER_BLOCK, vlen><<<grid_dim, block_dim>>>(in_ptr, out_ptr, num_ele);
在不同硬件上可微调 case 规模的阈值以选取不同实现策略(大规模走两阶段 ReduceLarge,小规模走单 Kernel 的 ReduceSmall)。
__threadfence() 保障所有 block 已完成计算且局部结果已写回全局内存,之后用其中一个 block 完成最终合并。原书此处标注"省略部分过程代码",最终块归约的具体实现细节、semaphores 数组初始化方式教材未给出完整代码。本课同样不做猜测性补全,只需理解 __threadfence() 保证"写入数据对其他 block 可见",atomicAdd 返回值判断"是否为最后一个完成的 block"是这种单 Kernel Grid 级同步的经典写法,实际使用需自行补全并在真实硬件上验证正确性。| 优化手段 | 解决的问题 | 关键代价 |
|---|---|---|
| 共享内存分治 | 减少全局内存频繁访问 | 需要同步 __syncthreads() |
| 减少 Warp Divergence(stride 从大到小) | 避免每一步都产生分歧 | 无 |
| Shuffle 指令 | 进一步减少共享内存访问,用寄存器直接线程间通信 | 依赖硬件是否支持 shuffle |
| Vector Load | 减少 block 数量、减少指令条数 | 需要数据对齐、寄存器占用增加 |
| 单 Kernel(Grid 级同步) | 避免二次 Kernel 启动开销 | 实现复杂,依赖原子操作与内存栅栏,正确性需仔细验证 |
案例 2:GEMV 优化
问题定义:GEMV(General Matrix-Vector Multiplication,通用矩阵-向量乘法):
\[ \text{out}[i] = \sum_{j=0}^{k-1} \text{mat}[i,j] \cdot v[j] \]将矩阵 mat 的每一行与向量 v 做点积,生成长度为 m 的输出向量。若考虑 float 类型数据,其算术强度为:
GEMV 的算存比通常远低于 1,因此是典型的内存受限(Memory Bound)工作负载,性能受限于内存带宽而非计算能力。
Step 1:一个线程束计算一个输出元素。每个线程束负责一个 out[i]:每个线程每次从全局内存加载 mat 和 v 对应位置的元素做乘加运算,每个线程累加部分 k 元素,最后在线程束内 Reduce(可用共享存储器,或若硬件支持 shfl 指令则用 __shfl_down_sync 做 Warp Shuffle Reduce)。若 k 较大,也可多个线程束协作计算一个输出点,借助共享存储器做跨线程束的 Reduce:
// 每个 warp 负责一行输出(依据教材文字描述重构的安全示例)
__global__ void gemv_v1(const float* mat, const float* v, float* out, int m, int k) {
int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
int lane_id = threadIdx.x % warpSize;
if (warp_id >= m) return;
float sum = 0.0f;
for (int j = lane_id; j < k; j += warpSize) {
sum += mat[warp_id * k + j] * v[j];
}
// Warp 内 Reduce(shuffle 版本)
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
if (lane_id == 0) {
out[warp_id] = sum;
}
}
Step 2:使用 128 比特 load 指令提高带宽效率。Step 1 中每个线程从全局内存加载 32 比特(单个 float),这种粒度通常不是最优的,会导致较多访存指令和低效带宽利用。将加载方式改为 128 比特(float4)可减少访存指令数量,缓解内存子系统压力,提高整体访存效率。
Step 3:用共享存储器缓存向量 v。线程块内部,mat 每一行仅对应输出元素的线程访问,而向量 v 需要被多个线程重复读取。可将 v 预先加载到共享存储器,避免每个线程都从全局内存访问相同数据。向量 v[k] 为线程块内所有线程共享,可分段并缓存到共享存储器,每次循环加载一段 v 再计算对应的 mat · v 子块。
关键调优参数 TILE_V:取值(如 512、1024 或 2048)直接决定算法的访存特征。增大 TILE_V 可提升数据复用率,但会显著增加单个 Thread Block 对 Shared Memory 的占用,进而降低硬件 Occupancy,需通过调优该参数寻找最优平衡点;Shared Memory 缓存向量还会引入额外一次同步操作,最终性能提升幅度需结合具体硬件的访存带宽与同步吞吐量实测评估。
Batch GEMV 的应用场景:尽管 GEMM 是广泛应用的高性能矩阵计算方案,但在某些场景(例如 [m, k]×[k, 4] 这类小尺寸 GEMM)难以充分发挥 GPU 的并行计算能力与带宽潜力,此时将计算转换为 Batch GEMV 形式往往能更有效利用内存带宽和缓存资源。判断何时该用 GEMM、何时该退化为 Batch GEMV,取决于矩阵的"瘦"或"胖"程度(N 维度是否过小导致 GEMM 分块策略无法有效利用硬件)。
案例 3:Flash Attention 优化
标准 Self-Attention 公式:
\[ O = \text{softmax}(QK^\top)V \]其中 \(Q, K, V, O\) 都是 2D 矩阵,shape 为 \((N, d)\),\(N\) 代表序列长度 seqlen,\(d\) 代表每个注意力头的维度 headdim。
\[ S = QK^\top, (N\times N) \qquad P = \text{softmax}(S), (N\times N) \qquad O = PV, (N\times d) \]通过调用高性能实现的 GEMM → Softmax → GEMM 流程,就可以完成标准 Self-Attention 计算。核心问题:\(S\) 和 \(P\) 作为中间结果矩阵,形状为 \((N,N)\),对显存的需求是 \(O(N^2)\),长 seqlen 时很容易爆显存,同时对显存的访问压力也会急剧增大。
Flash Attention 是一种更高效的算法,由斯坦福大学团队于 2022 年提出,旨在提升注意力机制的计算效率和访存效率,解决长序列的内存瓶颈问题。核心创新:① 分块计算——将大矩阵分解为小块进行计算;② 内存访问高效——通过重新计算(Recompute)避免存储完整的注意力矩阵;③ I/O 感知设计——优化 GPU 内存层次结构中的数据流动。
Online-Softmax:打破全局依赖
Softmax 计算依赖一个全局分母项,优化分块计算的关键是解决 Softmax 计算中的全局依赖问题,Online-Softmax 就是解决该问题的核心算法。为避免数值溢出,工程实现中采用 Safe-Softmax:
\[ \text{softmax}(x_1,\cdots,x_n) = \left\{ \frac{e^{x_i-m}}{\sum_{j=1}^{N} e^{x_j-m}} \right\}_{i=1}^{N}, \qquad m = \max_{j=1}^{N}(x_j) \]工程上朴素的三遍循环算法需要对序列 \(N\) 执行三次循环,对访存来讲相对低效,且由于第二步计算依赖第一步循环结束后的结果,无法简单融合。
Online-Softmax 的递推推导:构造一个与全局最大值 \(m_N\) 无关的 \(d_i'\):
\[ d_i' := \sum_{j=1}^{i} e^{x_j - m_i} \]当 \(i=N\) 时 \(d_N = d_N'\)。进一步推导 \(d_i'\) 和 \(d_{i-1}'\) 的递归关系:
\[ d_i' := d_{i-1}' e^{m_{i-1}-m_i} + e^{x_i-m_i} \]可以看到,\(d_i'\) 和 \(d_{i-1}'\) 之间的递归关系只依赖 \(m_i\) 和 \(m_{i-1}\),于是可以把求最大值的计算和 \(d_i'\) 的计算放在同一轮循环里,得到两遍循环的 Online-Softmax:
for i ← 1, N do
m_i ← max(m_{i-1}, x_i)
d_i' ← d_{i-1}' e^{m_{i-1}-m_i} + e^{x_i-m_i}
end
for i ← 1, N do
a_i ← e^{x_i - m_N} / d_N'
end
这种算法改造打破了 Softmax 对全局数据的依赖,使计算可以在分块(Tiling)流转过程中同步完成。通过 Online-Softmax,可以将 \(Q\)、\(K\) 矩阵计算的中间结果始终保留在 SRAM 中累加,无须将巨大的中间注意力矩阵写回显存——这种以计算换访存的策略,正是 Memory Efficient Attention 的核心原理。
合并为单轮循环
结合 Online-Softmax 后,Attention 原本三步计算缩减为两步:
for i ← 1, N do
x_i ← QK^T[:, i]
m_i ← max(m_{i-1}, x_i)
d_i' ← d_{i-1}' e^{m_{i-1}-m_i} + e^{x_i-m_i}
end
for i ← 1, N do
a_i ← e^{x_i - m_N} / d_N'
o_i ← o_{i-1} + a_i V
end
进一步观察到 \(o_i\) 的计算也可以不依赖 \(m_N\),构造 \(o_i'\) 后得到最终递推关系:
\[ o_i' \leftarrow o_{i-1}' \frac{d_{i-1}'}{d_i'} e^{m_{i-1}-m_i} + \frac{e^{x_i-m_i}}{d_i'} V[i,:] \]于是可以得到单轮循环的完整 Attention 算法:
for i ← 1, N do
x_i ← QK^T[:, i]
m_i ← max(m_{i-1}, x_i)
d_i' ← d_{i-1}' e^{m_{i-1}-m_i} + e^{x_i-m_i}
o_i' ← o_{i-1}' (d_{i-1}'/d_i') e^{m_{i-1}-m_i} + (e^{x_i-m_i}/d_i') V[i,:]
end
Flash Attention 实现:分块与 SRAM 调度
根据上面的推导,可以在一轮循环中计算完整的 Attention:对 \(Q, K, V\) 分块,将分块加载到高速 SRAM 中,在 SRAM 中完成当前线程块的 Attention 计算,最终将结果 \(O\) 写回显存,避免保存中间结果矩阵 \(P, S\)。在 Seq_Len_Q 上增加并行,一个线程块中只需一次 Q 加载;内循环按 Seq_Len_KV 进行多次 K、V 加载,这也是目前常用的 FA V2/V3 的切分方式。
确定块大小:Flash Attention 算法提出了 \(B_r\)(切分 Seq_Len_Q)和 \(B_c\)(切分 Seq_Len_KV)两个块大小概念,主要受限于硬件 SRAM 大小(对 MTT S5000 架构而言,指共享存储器):
\[ \text{SRAM}(Q) = B_r \times d \qquad \text{SRAM}(S) = B_r \times B_c \qquad \text{SRAM}(K,V) = 2 \times B_c \times d \]MTT S5000 具体配置示例:对于常见的 Headdim 128,在 192KB 的共享存储器限制下,\(B_r\) 选择 256,\(B_c\) 选择 128,可以刚好用满。由于 \(Q\) 一直缓存在共享存储器中,主循环内只有 \(K, V\) 的访存,不会成为明显瓶颈。一般而言倾向于缓存更多的 \(Q\)(尽可能大的 \(B_r\))以获得更高的计算访存比。矩阵运算由专用单元 TCE(Tensor Compute Engine)承载,其 \(A, B\) 输入数据来自 Shared Memory,输出 \(C\) 存放在寄存器中。一条 TCE 指令由一个 Squad(4 Warps) 协同执行:若将 ThreadBlock 配置为 512 线程(4 个 Squads),每个 Squad 在 \(B_r\) 方向均分后恰好对应 64 行(256/4),契合 TCE 原生支持的指令 Shape,因此单指令 ≥ M64 更有利于 TCE 的性能发挥。BlockSize 另一限制因素是寄存器数量:在 \(B_r=256, B_c=128\) 分块下,两次矩阵乘的 \(C\) 矩阵占据 128 个寄存器,处于 MTT S5000 硬件架构的承载范围内。
软件流水设计
Flash Attention 软件流水设计采用 TME 进行 \(Q, K, V\) 的数据复制操作。Attention 计算中 GEMM(QK)、Softmax、GEMM(PV) 在一轮循环里有顺序依赖关系,但 Softmax 中存在许多吞吐量较低的非 GEMM 操作(如 EXP2 计算)。为充分掩盖 Softmax 计算,可预先计算一轮 \(QK\),在主循环内让 GEMM(QK)_i 与 Softmax_{i-1} 重叠操作,实现更高硬件利用率:
PROLOGUE:
TME LOAD Q
TME LOAD K_0
TCE SQMMA Q K_0
MAIN LOOP i in (1, seq_kv // BN):
[并行执行]
├─ TME LOAD K_{i+1} (i>1) │ TME LOAD V_{i-1}
├─ TME LOAD V_i (i==1) │ TCE SQMMA Q K_i
└─ TCE SQMMA P V_{i-1} │ Softmax_{i-1}
LAST LOOP:
TCE SQMMA P V_n
Softmax_n
但这种软件流水的深度重叠需要在寄存器中额外存储一份 \(P\),需评估硬件寄存器数量是否够用。如果寄存器压力过大,容易导致编译器生成 Kernel 时产生寄存器溢出,显著降低性能。对 \(P\) 寄存器数量影响最大的可调节参数是 \(B_r\),减小 \(B_r\) 可避免寄存器溢出,但这与"尽可能选择大的块大小"相冲突,实践中应根据性能分析结果权衡。
TILE_V 的辩证关系是同构的优化矛盾——增大分块提升数据复用率,但会挤占有限的片上资源(寄存器/共享内存),最终都要通过实测寻找平衡点。建议把这三处联系起来理解,形成"资源图"式的整体认识。案例 4:GEMM 优化
问题定义:GEMM 是线性代数中最基本的操作之一,广泛用于深度学习、科学计算和图形处理,特别是神经网络的全连接层和卷积层。输入矩阵 \(A\) 是 \(M\times K\),\(B\) 是 \(K\times N\),输出 \(C\) 是 \(M\times N\):
\[ C_{ij} = \sum_{k=0}^{K-1} A_{ik} \times B_{kj} \]1. GEMM 初级实现(C 语言分块版)
分块不能降低计算量,目的是把局部数据缓存在高速存储器上。缓存矩阵 \(C\) 的数据后,进行一次乘累加计算,只需要对 \(C\) 进行一次读写、对 \(A, B\) 进行一次读。乘累加计算指令需要的数据都来自寄存器,所以 \(A, B\) 也需缓存,消耗的缓存量为 total_regs = tile_m×tile_n + tile_m×tile_k + tile_n×tile_k。如果 total_regs 一定,希望最大程度缓存 \(C\),线程的 tile_k 得为 1。
constexpr int tile_m = 8;
constexpr int tile_n = 8;
constexpr int tile_k = 8;
for (int m = 0; m < M; m += tile_m) {
for (int n = 0; n < N; n += tile_n) {
for (int k = 0; k < K; k += tile_k) {
for (int i = 0; i < tile_m; i++) {
for (int j = 0; j < tile_n; j++) {
for (int kk = 0; kk < tile_k; kk++) {
C[m + i][n + j] += A[m + i][k + kk] * B[k + kk][n + j];
}
}
}
}
}
}
n),推测为原书排版/印刷错误。上方代码已按正确写法 n += tile_n 更正呈现,教学时应提醒学员这是教材原文的疑似笔误,实际编写代码时应以此更正版本为准。MUSA naive 实现(对应线程内 tile_m=tile_n=tile_k=1),可直接编译运行:
/*
row_major, shape(row, col)
output(m, n)
input_a(m, k)
input_b(k, n)
*/
__global__ void mat_mul_naive(float *output, const float *input_a, const float *input_b,
int m, int n, int k) {
int tx = blockIdx.x * blockDim.x + threadIdx.x;
int ty = blockIdx.y * blockDim.y + threadIdx.y;
if (ty < m && tx < n) {
float c = 0;
for (int i = 0; i < k; ++i) {
c += input_a[ty * k + i] * input_b[i * n + tx];
}
output[ty * n + tx] = c;
}
}
这个版本分块小、并行度高,但计算访存比过低。并行程序性能上限由关键路径决定,过低的计算密度导致数据加载太频繁,计算单元长时间空闲等待数据导致算力浪费。要获得 GEMM 高性能实现,需要先对性能进行建模。
2. GEMM 性能建模
GEMM 计算量为 \(M\times N\times K\times 2\)(乘加算两次),双倍数据速率(DDR)读取访存量为 \((M\times K + N\times K)\times d_{type\_size}\),计算访存比(只考虑读取)为:
\[ \frac{M\times N\times 2}{(M+N)\times d_{type\_size}} \]公式中的 \(M, N, K\) 指的是并行执行分块的 \(M, N, K\),而非 GEMM 问题参数本身——在 DDR 层面可能并行执行的数值有几千,MP 层面只有几个,线程层面是个位数,计算访存比这个概念是分层的。
关键工程结论:① 高端 GPU 计算吞吐量为几百 TOPS,但 DDR 吞吐量为几 TB/s,如果要跑满计算吞吐量,计算访存比至少要达到 100,即每字节 DDR 加载至少要支撑 100 次运算;② 从公式看计算密度过低会归为 Memory Bound,benchmark 应考察内存带宽达成率而非计算吞吐量达成率;③ \(K\) 分块的选择不宜过大,尽量让片上空间用于 \(M, N\) 分块的扩大;④ 在 \(M\) 和 \(N\) 乘积一定的情况下,\(M==N\) 时计算访存比最高,所以并行计算的分块最好是一个方块。
寄存器带宽约束:以寄存器为例,一条 FMA(乘累加)指令需要读 3 个寄存器、写 1 个寄存器,如果通用寄存器读取带宽只能达到每周期 1 个寄存器,那么 GEMM 峰值性能只能达到浮点峰值的 33%。如果每个线程计算 8×8 分块,发射 64 条 FMA 指令需要 64 cycle,共享存储器需读 8 个 A 寄存器和 8 个 B 寄存器(1 个寄存器 4 字节),共享存储器到寄存器带宽需求为每 64cycle 读 64 字节,即每线程每 cycle 读取 1 字节;考虑 MT GPU 每个 MP 每周期有 128 个线程被激活,MP 需提供 128 字节/cycle 的寄存器到共享存储器带宽——这在没有存储体冲突的情况下刚好满足需求。若有效率折损 50% 的存储体冲突,最终 GEMM 计算吞吐量上限为浮点计算峰值的 50%。
3. GEMM 高性能实现(伪代码结构)
mat_mul_pseudo 是教材明确标注为伪代码的示例,其中 load_gmem_to_reg、store_reg_to_smem_with_transpose、mma8x8、swizzle_blkid 等均为未定义的辅助函数,用于表达算法结构,不能直接编译运行。实际工程实现需要结合具体硬件 ISA 补全这些函数,本课不做猜测性补全。__global__ void mat_mul_pseudo(float *output, const float *input_a, const float *input_b,
int m, int n, int k) {
constexpr int tile_size = 128;
constexpr int TILE_K = 4;
// ping pong switch is optional
// shared memory, 2row for double buffer
__shared__ float __attribute__((aligned(16))) smem_a[2][TILE_K * tile_size];
__shared__ float __attribute__((aligned(16))) smem_b[2][TILE_K * tile_size];
// register
float4 load_reg_a[2];
float4 load_reg_b[2];
float4 reg_a[2][2]; // 2row for double buffer
float4 reg_b[2][2];
float4 reg_c[8][2] = {{float4(0.f, 0.f, 0.f, 0.f)}};
swizzle_blkid(blockIdx.x, blockIdx.y, blockIdx.z);
// offset index compute(省略:需按实际布局计算)
int load_offset_a = 0;
int load_offset_b = 0;
int store_offset_a = 0;
int store_offset_b = 0;
// load first block from global memory to shared memory
load_gmem_to_reg(input_a, load_offset_a, load_reg_a);
load_gmem_to_reg(input_b, load_offset_b, load_reg_b);
store_reg_to_smem_with_transpose(load_reg_a, store_offset_a, smem_a[0]);
store_reg_to_smem(load_reg_b, store_offset_b, smem_b[0]);
__syncthreads();
// load first register from shared memory
load_smem_to_reg(smem_a[0], 0, reg_a[0]);
load_smem_to_reg(smem_b[0], 0, reg_b[0]);
for (int i = TILE_K; i < k + TILE_K; i += TILE_K) {
load_gmem_to_reg(input_a, load_offset_a + i, load_reg_a);
load_gmem_to_reg(input_b, load_offset_b + i, load_reg_b);
for (int j = 0; j < TILE_K - 1; j++) {
load_smem_to_reg(smem_a[(i - TILE_K)/TILE_K], j + 1, reg_a[(j + 1) % 2]);
load_smem_to_reg(smem_b[(i - TILE_K)/TILE_K], j + 1, reg_b[(j + 1) % 2]);
mma8x8(reg_c, reg_a[j % 2], reg_b[j % 2]);
}
if (i < k) {
__syncthreads();
store_reg_to_smem_with_transpose(load_reg_a, store_offset_a, smem_a[i/TILE_K]);
store_reg_to_smem(load_reg_b, store_offset_b, smem_b[i/TILE_K]);
__syncthreads();
}
load_smem_to_reg(smem_a[i/TILE_K], 0, reg_a[0]);
load_smem_to_reg(smem_b[i/TILE_K], 0, reg_b[0]);
mma8x8(reg_c, reg_a[1], reg_b[1]);
}
store_reg_to_gmem(output, reg_c);
}
优化策略详解(对照伪代码逐项走查):
(1) 线程块与计算任务分配:每个线程块包含 256 个线程,负责计算一个 128×128 的矩阵块,即每个线程计算 8×8 个矩阵元素,高效利用 GPU 并行计算能力,避免线程间资源竞争。
(2) 共享内存到寄存器——双缓冲(Double Buffer):smem_a 和 smem_b 在 K 维度上设置双缓冲,计算循环中一部分共享内存块用于当前计算、另一部分用于加载下一轮数据。每次从共享内存搬运 16B 数据,K 方向每次搬运 1 个数,配合适当 Padding 可避免 bank 冲突,双缓冲有效掩盖了内存访问延迟。
(3) 全局内存到共享内存——双缓冲:每个线程从全局内存加载连续位置的数据,通过寄存器缓存后存储到共享内存,每次搬运 16B 数据,K 方向每次搬运 4 个数,保证内存带宽高效使用;相邻线程访问连续内存位置,提升带宽利用率,避免不连续访问导致的性能瓶颈。
(4) Block ID Swizzling——提高 Cache 命中率:通过 swizzle_blkid 函数重新排列线程块 ID,使线程块能均匀访问不同内存区域,减少不同线程块对相同内存位置的访问冲突,提高缓存命中率、减少全局内存访问频率。若是 Tensor Core 的 GEMM 实现,可以把伪代码中的 mma_8x8 换成对应的 tensor core mma 指令,同时保证输入输出数据 layout 符合 mma 指令定义。
Block ID Swizzle 深度解析
① 读共享存储器存储体冲突的解决方法(3 种):若写带宽需求较低,可规划无冲突的读布局让写共享存储器承受冲突的副作用;若共享存储器空间够用,可在布局中加入 pad;也可使用 Swizzle 的共享存储器布局。② 寄存器的存储体冲突:汇编层面要尽量保证一条指令中的寄存器分属不同存储体,C 语言程序则要靠编译器优化解决。
③ Block ID 的 Swizzle 几何解释:GPU 默认 block ID 执行顺序是 blockIdx.x → blockIdx.y → blockIdx.z。假设 \(M\) 和 \(N\) 足够大,GPU 上同时刻并行执行的 block 数量大概率与 blockIdx.y 是同一个值、blockIdx.x 递增。若 blockIdx.y 映射到 \(M\)、blockIdx.x 映射到 \(N\),矩阵 \(A\) 被不同 MP 共享而矩阵 \(B\) 无法被共享,在 \(M\) 和 \(N\) 分块一样大的情况下,缓存最多只提供接近两倍的 DDR 带宽放大效果。
Swizzle 后的效果:如果每 8 个 block 换一行,只需要 8 个 tile_m 的数据和 8 个 tile_n 的数据就能支持 64 个 block 一轮的计算,64 个 block 提供的 DDR 带宽放大效果为 (64+64)/(8+8) = 8 倍,所以让 blockIdx 的遍历提前换行可以让缓存的带宽放大效果更明显:
未 Swizzle(一字排开):
(0,0) (0,1) (0,2) (0,3) ... (0,60) (0,61) (0,62) (0,63)
→ 矩阵 A 复用,矩阵 B 几乎不复用,放大比约 2 倍
Swizzle 后(每 8 个 block 换一行):
(0,0) (0,1) (0,2) (0,3) (0,4) (0,5) (0,6) (0,7)
(1,0) (1,1) (1,2) (1,3) (1,4) (1,5) (1,6) (1,7)
...
(7,0) (7,1) (7,2) (7,3) (7,4) (7,5) (7,6) (7,7)
→ 矩阵 A、B 都在局部得到复用,放大比约 8 倍
Swizzle 的 Python 参考实现(教材完整给出,可直接在 Python 环境执行):
def swizzle_naive(nr_tile, macro_tile_x, macro_tile_size, res_tile_x, macro_bid_max,
max_tile_y, switch_flag, grid_y, grid_x):
valid_set = set()
rst = []
for blockidx_xyz in range(nr_tile):
macro_x, macro_res = blockidx_xyz // macro_tile_size, blockidx_xyz % macro_tile_size
if blockidx_xyz >= macro_bid_max:
bidy, mini_x = macro_res // res_tile_x, macro_res % res_tile_x
if bidy % 2 == 1:
mini_x = res_tile_x - 1 - mini_x
else:
bidy, mini_x = macro_res // macro_tile_x, macro_res % macro_tile_x
if bidy % 2 == 1:
mini_x = macro_tile_x - 1 - mini_x
if macro_x % 2 == 1:
bidy = max_tile_y - bidy
bidx = macro_x * macro_tile_x + mini_x
if switch_flag:
bidx, bidy = bidy, bidx
idx = (bidy, bidx)
rst.append(idx)
assert idx not in valid_set
assert bidx < grid_x and bidy < grid_y, f'fail at {blockidx_xyz} with {idx}'
valid_set.add(idx)
print(len(valid_set), len(rst))
return rst
# 参数初始化与内核启动模拟
switch_flag = False
grid_x = 10
grid_y = 10
nr_tile = grid_x * grid_y
macro_tile_x = 8
if not switch_flag:
macro_tile_size = macro_tile_x * grid_y
res_tile_x = grid_x % macro_tile_x
macro_bid_max = grid_x // macro_tile_x * macro_tile_x * grid_y
max_tile_y = grid_y - 1
else:
macro_tile_size = macro_tile_x * grid_x
res_tile_x = grid_y % macro_tile_x
macro_bid_max = grid_y // macro_tile_x * macro_tile_x * grid_x
max_tile_y = grid_x - 1
print(swizzle_naive(nr_tile, macro_tile_x, macro_tile_size, res_tile_x, macro_bid_max,
max_tile_y, switch_flag, grid_y, grid_x))
该函数处理了"蛇形折返"(bidy % 2 == 1 时翻转 mini_x;macro_x % 2 == 1 时翻转 bidy)与"边界不整除时的余量处理",是一个较完整的工程级 Swizzle 实现,建议用 grid_x=grid_y=10, macro_tile_x=8 的具体参数跑一遍,观察输出的 (bidy, bidx) 序列如何从"逐行遍历"变成"分块折返遍历"。
④ 访存合并(收束):伪代码中用 float4 访问,可保证相邻线程访问连续地址空间,降低相同传输任务量下存储系统的负载。可以把 MP、各级缓存、DDR 看成独立主机,这些主机间通过网络报文通信,片上各种总线充当交换机——报文越长,相同访存量下的报文数量就越少,交换机压力越低,各级主机处理报文的效率也更高,整个网络的拥塞程度会降低,访存的平均延迟和最差延迟表现都有保障。
swizzle_naive Python 代码(grid_x=grid_y=10, macro_tile_x=8),打印并对比 Swizzle 前后 block 遍历顺序。④ 讨论题:为什么"缓存矩阵 C、让 tile_k=1"是在 total_regs 一定时的最优选择?六、版本与兼容性
| 项目 | 本课校准基线 | 备注 |
|---|---|---|
| MUSA SDK | 5.2.0 | 本课涉及的内存优化、多卡通信、性能优化 API 以此版本头文件与官方文档为准 |
| torch_musa | 2.9.1 | 如在 PyTorch 生态下复现 Reduce/GEMV/FlashAttention/GEMM 案例,以该版本行为为准 |
| 参考 GPU | MTT S5000(平湖) | Occupancy 数值示例、共享内存容量(192KB)、TCE/Squad 描述、1024 比特向量访存等均以此型号为基准,其他型号需按各自规格重新计算 |
| 通信库 | MCCL(多卡 all-reduce/broadcast) | 版本号原书未给出,以当前 MUSA SDK 配套版本为准 |
| 代码可运行性 | 见各代码块前后说明 | Reduce(ReduceLarge/ReduceSmall/WarpReduce)与 GEMM naive(mat_mul_naive)标注为可编译运行;GEMV(gemv_v1)、P2P 调用示例为依据文字描述的安全重构;GEMM 高性能实现(mat_mul_pseudo)为教材明确标注的伪代码,不可直接编译 |
| 存疑项清单 | 不做猜测性补全 | ① Bank Conflict 小节 fast_swizzle 的 mask/shift 具体取值;② musaMemcpyPeer/musaDeviceEnablePeerAccess 完整调用代码(教材仅文字描述,本课已重构为安全示例并标注);③ Reduce 单 Kernel 扩展方案(ReduceInOne)中 Grid 级同步的完整代码(教材注明"省略部分过程代码");④ GEMV 完整实现(GEMV.cpp 不在本次提取页面范围内);⑤ Flash Attention 软件流水示意图为依据文字描述的重构呈现,非原书图 8-12 的精确复刻;⑥ 表 8-1 各级存储器供需关系的具体数值(原书扫描件复杂多列表格,识别存在不确定性);⑦ GEMM 初级实现代码中的疑似笔误已按逻辑更正,仅供参考 |
musaGetDeviceCount、musaSetDevice、musaMemcpyPeer、musaDeviceEnablePeerAccess、musaStream_t 等)与 __shfl_down_sync 等内建函数,均可能随 SDK 版本迭代发生签名、行为或参数默认值的变化。务必以本机已安装的 MUSA SDK 5.2.0 头文件与当前官方文档为准,发现纸质教材描述与实际编译报错不一致时,优先信任头文件与官方文档。七、本课与相邻课程
- m11-l9(第 7 章:MUSA 程序调试):本课是第 7 章调试技能的直接延续——性能优化前必须先保证程序正确,若 Reduce/GEMM 等 Kernel 结果不对,应先用 musa-gdb、Compute Sanitizer 等第 7 章工具定位正确性问题,再进入本课的性能调优流程。
- m11-l8(第 6 章):第 6 章介绍的 MUSA 存储器模型(全局内存、共享内存、寄存器层次)是理解本课 Occupancy、合并访存、Bank Conflict 的前置基础,建议先掌握存储器层次再学习本课。
- m11-l4(第 4 章:muBLAS/muDNN):本课手写的 GEMM/GEMV 高性能实现,与调用 muBLAS/muDNN 高性能算子库是"知其所以然"与"拿来即用"的关系——工程实践中优先调用成熟算子库,本课的手写实现有助于理解库内部的优化原理,也用于库无法覆盖的定制场景。
- m11-l10(第 9 章:性能分析工具):本课涉及的 Occupancy、访存带宽利用率、寄存器压力等指标,在实际工程中应结合第 9 章的性能分析工具(Profiler)实测验证,而非仅凭理论公式估算。