HSTU attention development and optimization using Cutlass/CuTe

Gems Guo, Jerry Chen, Bin Chai, NVIDIA DevTech Engineer | 2025.05.30

目录 (Agenda)

• HSTU 模块介绍
• Attention 优化
• Rab & dRab Cute 实践
• 问答环节

HSTU 模块介绍

模型架构概览

该模型是一种用于生成式推荐的序列转换器。其基本架构从原始特征（数值和类别特征）开始，经过嵌入层和底层神经网络，进入特征交互神经网络（如 FMs, DCN, Transformers, DHEN 等），最终到达顶层神经网络（如 MMoE, PLE 等）。

<center>图注：左侧为典型的推荐系统神经网络分层结构。右侧为序列化统一特征处理流程，其中绿色高亮部分是本次优化的核心计算模块。</center>
<center>引用自论文: Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations</center>

HSTU (Multi-Head Attention) 模块

HSTU 是一个多头注意力模块，作为 Transformer 的核心组件。它包含在一个残差连接和层归一化（Add & Norm）的结构中，并可重复 N 次。该模块内部的具体流程为：线性变换 -> 分割 -> 注意力计算 -> 拼接 -> 线性变换。其中，相对位置偏置（Relative Attention Bias）被引入到注意力计算中。

性能瓶颈：内存占用

在实际应用中，中间张量（尤其是 QKᵀ）会占用巨大的显存空间，成为性能瓶颈。

假设参数如下：
- num_blocks = 16
- batch_size = 128
- sequence_length = 1000
- num_heads = 8
- dim_x = 256, dim_qk = 32, dim_uv = 32

各张量的维度和内存占用估算如下：
- X, Y: sequence_length × dim_x
- Q, K: sequence_length × num_heads × dim_qk
- U, V: sequence_length × num_heads × dim_uv
- QKᵀ: sequence_length² × num_heads
- AV: sequence_length × num_heads × dim_uv

此时，内存占用估算：
- X, Y 的内存大小 ≈ 0.98 GiB
- U, V, Q, K 的内存大小 ≈ 0.98 GiB
- QKᵀ 的内存大小 ≈ 30.52 GiB
- AV 的内存大小 ≈ 0.98 GiB

巨大的 QKᵀ 张量（约 30.52 GiB）是主要的内存瓶颈。

Attention 优化

解决方案：算子融合 (Kernel Fusion)

为了解决上述内存瓶颈，核心思想是进行算子融合，避免将巨大的中间结果 QKᵀ 写入全局内存。

• 前向传播 (Forward Pass)：可以将多个计算步骤融合成少数几个 CUDA Kernel。QKᵀ 的计算结果直接存储在寄存器（registers）中作为中间结果，供后续计算使用，而无需写入显存。
• 反向传播 (Backward Pass)：同样地，反向传播中的计算也可以被融合成少数几个 Kernel，并将中间结果存储在寄存器中。

前向与反向传播计算图

下图详细展示了前向传播（左）和反向传播（右）的数据流图。这些计算步骤是进行算子融合的候选对象。

• 前向传播：输入 Q, K, V, U，经过 SiLU 激活，然后是一系列的矩阵乘法（MatMul）、加法（Add Relative Attention Bias）、Masked Fill、SiLU、Norm 和 Hadamard 积，最终得到输出 O。
• 反向传播：根据输出梯度 dO，计算各个输入的梯度 dU, dV, dQ, dK 以及其他中间变量的梯度。

反向传播中的重计算

在反向传播过程中，梯度的计算依赖于前向传播中的中间结果 QKᵀ。为了避免在前向传播时存储这个巨大的张量，可以在反向传播时重新计算 (recalculated) QKᵀ，并将其结果保存在寄存器中供梯度计算使用。这是一种典型的用计算换空间的优化策略。

层归一化 (Layer Normalization) 的反向传播优化

层归一化（Layer Normalization）的反向传播有多种等价的计算公式。不同的公式在计算效率和内存访问上有不同的权衡。下表对比了四种计算 dxᵢ 的等价公式的开销。选择最优的公式对于融合后的 Kernel 性能至关重要。

计算规模与指令分析

以一个具体实例为例，历史行为序列长度为 2048，维度为 256 x 18。注意力矩阵的计算如下图所示，其中黄色部分表示由于因果依赖（causal dependency）而被屏蔽的部分。

下表展示了该任务中各种指令的理论计数和实际测量计数，量化了计算和内存访问的负载。

• Matrix Multiply and Accumulate: 理论约 9,510,912 次，实测 10,027,008 次。
• Load Matrix from Shared Memory: 理论范围 811,008 至 7,281,792 次，实测 5,796,864 次。

优化原则

基于以上分析，可以总结出以下四个核心优化原则：

1. 利用高速片上内存 (如 SMEM 或 REG) 来缓存数据，从而减少内存访问指令的数量并减轻各级内存的压力。
2. 调整计算流 (指令顺序) 和同步时机，使得计算和内存访问可以相互重叠，从而隐藏内存访问延迟。
3. 展开计算或内存访问指令，以利用指令间并行性，提高流水线利用率。
4. 减少非必要的指令数量。

待探索的优化方向

为了达到极致性能，可以在以下几个方面进行探索和优化：

• 探索不同的分块（tiling）大小和相关的计算工作流。
• 在线程块（thread block）内部，对线程束（warp）进行特化。
• 在整个数据范围内，对线程块进行特化。
• 在保持拓扑顺序的同时，搜索最优的计算序列和同步时机。
• 平衡流式多处理器（streaming multiprocessors）之间的负载。
• 探索用于复制和计算的最佳迭代粒度。
• 研究计算 SiLU、Filled Mask 操作的最优指令。
• 探索加载 Q/K/V/Bias 和存储 O 的最高效方式。

负载均衡挑战

注意力计算中存在非均匀的计算模式（如下图中的矩形和三角形区域），这给流式多处理器上的负载均衡带来了挑战。不同的计算区域可能需要不同的处理策略，以避免部分处理器空闲，从而最大化硬件利用率。

注意力矩阵的分块计算方法

通过对注意力矩阵采用分块（tiling）方法，内存访问和计算只在掩码矩阵（mask matrix）中有意义的“分块（tile）”上执行。这种方法有效地节省了不必要的计算能力和带宽。

HSTU 注意力核（Kernel）概览

• 性能远超 OAI-triton。
• 在 NVIDIA/RecSys-example 和 pytorch/FBGEMM 中开源。
• 支持应用于生成式推荐（Generative Recommendation）的各种定制化掩码（mask）。
• 支持用于训练的实验性 FP8 注意力机制。
• 支持用于推理的分页注意力（Paged Attention）。

下图展示了“上下文 + 组目标掩码（contextual + group target mask）”的结构。

优化概览

优化的整体流程如下图所示，涵盖了从核函数融合到掩码设计的多个步骤。

Fused Attention 前向传播

Fused Attention 的前向传播过程通过将多个操作（如 GEMM、加偏置、激活函数、掩码）融合成一个 CUDA 核来减少内存访问开销。

伪代码如下：

For i in Q_tile
  acc_o = {0}
  For j in K/V tile
    acc_s = GEMM(Q_tile[i], K_tile[j])
    acc_s += Bias_tile[i, j]
    acc_s *= alpha
    P = silu(P)
    P = mask(P)
    acc_o += GEMM(P, V_tile[j])
  acc_o *= scale

说明：
1. CTA Tile 沿 Seq_Q 维度划分。如果沿 Seq_K 维度划分，那么在计算输出 O 时会引入一次全局原子加（global atomicAdd）操作。

注意力机制中的分块大小（Tiling Size）

寄存器和共享内存的使用

在不同的分块大小和参数下，寄存器（register）和共享内存（shared memory）的使用情况如下：

• 寄存器使用:
• Dim 越大，所需的寄存器越多，因为 tOrV 和 acc_o 的大小与 DIM 相关。
• kBlockM 越大，所需的寄存器越多，因为 tSrQ 和 acc_s/acc_o 的大小与 kBlockM 相关。
• kBlockN 越大，所需的寄存器越多，因为 tSrK 和 acc_s 的大小与 kBlockN 相关。
• 共享内存使用:
• kBlockM/N 和 DIM 越大，所需的共享内存越多，因为分块大小与共享内存相关。
• 在 G2S（Global to Shared）加载 Q tile 后，我们可以通过 S2R（Shared to Register）将 Q tile 保留在寄存器中或留在共享内存中。
• S2R 操作会增加寄存器压力，但能减少共享内存压力并减少 LDS M 指令的数量。

不同分块大小下的指令数量

• 每个 CTA（Cooperative Thread Array）需要对完整的 K 和 V 执行 G2S 和 S2R 操作，这意味着 K 和 V 存在冗余的 G2S 和 S2R。
• 因此:
• BlockM 越小，会导致更多的块（block），从而为 K 和 V 带来更多冗余的 G2S 和 S2R 操作，进而产生更多的 LDS M/LDGSTS 指令。
• 如果 Q_tile 驻留在共享内存中，在进行 QK GEMM 时需要重复地对 Q_tile 执行 S2R 操作。这意味着 BlockN 越小，seq_K 方向上的分块数量就越多，从而导致 S2R 操作的重复次数增加以及更多的 LDS M 指令。

TileMMA 在注意力前向传播中的应用

Tile Size 对 TiledMMA 的影响

• 为 TileMMA 设置合适的 Tile Shape。
• 采用满足 MMA 指令的最小粒度。
• 尽可能使用 LDS M x4。

不同 Warp 排列对 TiledMMA 的影响

• 需要为 TileMMA 设置合适的 Warp 排列。
• 当 Seq_Q 较小时，通常建议使用右侧的方案，因为它会启动更多的块来充分利用 GPU。
• 当 Seq_Q 较大时，建议使用左侧的方案，因为它不会引入额外的块规约（Block Reduce）操作。

如何加载偏置（Bias）

LDG vs LDGSTS/LDSM

比较了两种加载 Bias 的方式：直接从全局内存加载（LDG）和通过共享内存加载（LDGSTS + LDSM）。

• LDG:
• 优点: 不增加共享内存压力。
• 缺点: 非合并的内存加载。
• LDGSTS + LDSM:
• 优点: 流水线可以被更精细地控制，而 LDG 依赖于编译器的优化。
• 缺点: 增加了共享内存的压力；更依赖于手动的流水线控制。

流水线（Pipeline）优化

Ampere 架构前向传播流水线优化

通过重新安排指令顺序，可以获得更好的流水线设计，提高指令的并行度。下图中，将 K_tile(0) 和 Bias_tile(0) 的加载操作提前，与 Q_tile 的加载并行，从而优化了 prologue 阶段。

在 Ampere 架构上让 SiLU 和 MMA 操作重叠

通过精心设计的 warp 级流水线，可以使 SiLU 激活函数的计算与 SV GEMM 操作重叠。此外，NVCC 编译器也会帮助进行指令重排，以实现指令级并行（ILP）。

如何移除 __syncthreads()

问题：在 mainloop 中，如果移除了 __syncthreads()，可能会出现读写冲突。例如，Warp 1 可能在 Warp 0 完成对共享内存中 Bias 数据的写入之前就开始读取，导致数据竞争。

解决方案：让每个 Warp 读写各自独立的共享内存区域。这样就不再需要线程同步。通过修改 GmemTiledCopyBias 中的一个参数，可以改变内存布局，实现这一目标。该优化可为 DIM=32/64 的情况带来 3-7% 的性能提升。

Ampere 架构前向传播针对小 DIM 的双缓冲（Double-buffer）流水线

问题：对于维度（DIM）较小的情况，计算无法完全隐藏数据拷贝的延迟，使得长计分板（long scoreboard）成为阻碍指令发出的主要瓶颈。
解决方案：设计一个双缓冲流水线，让数据拷贝“在飞行中”完成（make copy on the flight）。
该方案通过在 prologue 阶段预加载两份数据，使得 mainloop 在处理当前数据块的同时，可以预取下一个数据块，从而提升流水线效率。对于 DIM=32/64 的情况，性能可提升 8-15%。

Pipeline: Hopper Forward

遵循 FA3 设计，流水线方案为 Ping pong schedule + 2-stage HGMMA pipeline。

Ping pong schedule 流水线:
该图展示了两种流水线方案。
1. 上图 (2-stage HGMMA pipeline): 工作组1 (WG1) 和工作组2 (WG2) 交错执行 QK GEMM, PV GEMM 和 SiLU 操作。wait<1> 表示等待一个阶段。该方案中，WG1 执行 QK1 GEMM (1) 的同时，WG2 在执行 QK1 GEMM (0)。存在计算和等待的交替。
2. 下图 (ping pong schedule): 同样是两个工作组，通过 wait<0> 实现更紧密的流水线。WG1 执行 PV GEMM (0) 时，WG2 执行 QK1 GEMM (1)，实现了 HGMMA 和 SiLU 之间更好的重叠。

使用 2-stage HGMMA 的优缺点:
* 优点:
1. 可以发出更少的 wait 指令，从而发出更多的 WGMMA 指令，保持 MMA 流水线繁忙。
2. 2-stage WGMMA 流水线: issue QK HGMMA -> issue PV HGMMA -> wait<1> -> SiLU -> QK HGMMA -> ...
3. ping pong 流水线: issue QK HGMMA -> wait<0> -> SiLU -> issue PV HGMMA -> wait<0> -> issue QK HGMMA.
* 缺点:
1. 增加了寄存器压力，需要维护额外的 tSRs (tensor state registers)。

Pipeline: Ping pong schedule 实现

该页面展示了如何使用 cutlass::arch::NamedBarrier 来实现乒乓调度。

代码实现:
假设有3个工作组 (WG_nums = 3)。
* 初始化 Barrier: 在主循环之前，每个工作组 (WG) 到达（arrive）前一个 WG 的屏障。例如，WG2 arrive WG1 的屏障，WG1 arrive WG0 的屏障，WG0 arrive WG2 的屏障，形成一个环形依赖。
* 主循环 (mainloop): 在主循环中，每个 WG 首先同步并等待（sync）它自己的屏障，该屏障由前一个 WG 在上一轮迭代中释放。然后，它 arrive 下一个 WG 的屏障，以释放在下一轮迭代中等待的 WG。

状态表:
下方的表格展示了 Barrier 状态随时间步的变化：
* Step->: 初始状态，WG0 issue WGMMA，WG1 和 WG2 等待。
* Barrier init: WG0、WG1、WG2 分别 arrive 它们的前一个 WG 的屏障。
* Barrier status back to init: 经过一系列步骤后，状态返回初始状态，形成循环。

这个机制确保了工作组之间的同步，实现了高效的乒乓流水线调度。

Pipeline: Kernel Scheduler: Warp Specialized and Non-Persistent

该方案中，计算 prologue 和 epilogue 的延迟是暴露的。当主循环（mainloop）很短时，这个问题会更加明显。

• Producer: WG1: 负责执行 prologue 部分的工作。
• Consumer: WG1+WG2..: 负责执行 mainloop 和 epilogue。

这种非持久化的调度方式导致 prologue 和 epilogue 成为性能瓶颈，因为它们的执行时间没有被其他计算隐藏。

Pipeline: Kernel Scheduler: Warp Specialized with Dynamic Persistent

为了解决 Non-Persistent 调度中的延迟暴露问题，采用了动态持久化调度器，通过重叠不同 m_block 的 prologue 和 mainloop/epilogue 来隐藏延迟。

调度流程:
* Producer (WG1): 计算当前 m_block 的 Q(0), K, Bias，同时可以预取（prefetch）下一个 m_block 的数据 Q(1), K, Bias。
* Consumer (WG1+WG2..): 在 Producer 准备好当前 m_block 数据后，执行 PV GEMM。
* 重叠执行: 通过 ping-pong schedule，当 Consumer 处理当前 m_block 时，Producer 已经开始处理下一个 m_block 的 prologue，从而将 next m_block prologue 的延迟隐藏在当前 m_block 的 mainloop 和 epilogue 中。

代码实现:
* 使用 tile_id_counter (一个原子计数器) 来动态分配 tile 给不同的工作组。
* Producer: 循环获取 tile_id，执行 do_load，然后通过 smen_tile_id 和 sync(QueryEmpty) 通知 Consumer 数据准备就绪。
* Consumer: 等待 QueryEmpty 信号，获取 tile_id，执行 do_mma，然后通知 Producer QueryEmpty 信号，表示它可以接收新的数据了。

这种 producer-consumer 模式配合动态调度，有效地提升了流水线效率。

代数优化 (Algebraic Optimization)

通过代数变换来减少计算函数所需的指令数量。

双曲正切函数 (hyperbolic tangent function) 优化:
* tanh(x) 的标准形式可以变换，将 z 替换为 x/2，可以推导出 tanh(x/2) = 2 * sigmoid(x) - 1。
* tanh 的一个版本可以表示为 sigmoid(x) = 0.5 * tanh(0.5 * x) + 0.5。

SiLU 激活函数优化:
* 标准 SiLU (exp SiLU): SiLU = x * sigmoid(x)
* tanh SiLU: SiLU = x * (0.5 * tanh(0.5 * x) + 0.5)
* fast tanh SiLU: SiLU = x * tanhf(x) + x (通过手动简化)

指令数量对比:
下表展示了不同实现方案所需的指令：

Plan	Instruction
exp SiLU	NxFMA+NxFMUL+MUFU.RCP+MUFU.EX2+NxFADD
_expf SiLU	2xFMUL+MUFU.RCP+MUFU.EX2+FADD
tanhf SiLU	3xFMUL+MUFU.TANH+FFMA
fast _tanhf SiLU	FMULD2+MUFU.TANH+FFMA

通过采用单精度内部函数（Single Precision Intrinsics）可以增量地采用这些优化，以获得应用所需的精度。例如 fmaf, __expf, __tanf。fast math 可以更有效地减少计算数学函数的指令数。

工作负载均衡 (Workload balance)

块调度策略：最长处理时间优先 (Longest Processing Time first)

问题背景:
* 在因果掩码（casual mask）的情况下，每个任务（CTA）需要计算的 token 数量不相等，导致工作负载不均衡。
* 通常，CUDA 会按 blockIdx.x 从 0 到 gridDim.x - 1 的递增顺序分派 CTA。如果任务处理时间随 blockIdx.x 递减，则最后分派的短任务会很快完成，而先分派的长任务仍在运行，导致SM利用率下降，产生"拖尾效应" (tail effect)。

解决方案 (LPT First):
* 使用“最长处理时间优先”算法，首先启动具有较长工作负载的任务。
* 如右侧代码所示，通过修改 m_block 的计算方式，将 blockIdx.x 映射到网格的末端开始，实现了任务分派顺序的逆转。
* aggressive 模式: m_block = gridDim.z - blockIdx.z - 1
* 这种方式有助于缓解拖尾效应，从而在 SMs 之间实现更好的负载均衡。
* 如果严格遵循LPT算法，需要优先处理不同批次和头的任务，这会对L2缓存产生巨大压力。

工作负载均衡性能对比

下图展示了在 Ampere 架构上，不同块调度策略（Default, Moderate, Aggressive）在不同序列长度下的 HTSU Attention Kernel 延迟对比。

• 测试配置: head_dim=128, head_num=1
• 序列长度: seq_len=4096, seq_len=8192, seq_len=16384
• 结果分析:
  • 在所有测试的序列长度下，Aggressive 策略（即 LPT First）都表现出最低的延迟，性能最优。
  • 随着序列长度和批大小的增加，Aggressive 策略相较于 Default 策略的性能优势更加明显。例如，在 Batch Size 为 128 时，Aggressive 策略的归一化延迟为 100.0%，而 Default 策略为 127.9%。

这证明了 LPT First 调度策略在处理非均匀工作负载时的有效性。

如何进行掩码 (How to Mask)

掩码函数的几何化 (Geometrization of the mask function)

通过将掩码操作几何化，可以更高效地进行计算。

核心思想:
* 使用 Tiled_mma 获取 threadIdx 的分片。
* 通过 thr_mma.get_thread_slice(threadIdx) 和 thr_mma.partition_C(cs) 获取线程对应的张量分片坐标 tScS。
* 将局部坐标 (x, y) 转换为全局坐标，其中 x = int(get<0>(tScS(i))) + m_block * kBlockM;。
* 图示中，一个大的张量被划分为多个小的 Tile (V0, T1, ...)，右侧的 Tensor c5 展示了从全局张量中提取出的一个子张量的坐标。Tensor tScS 则是这个子张量中一个线程所负责处理的坐标子集。

掩码函数的几何化实现细节

实现流程:
1. 块边界常量 (Block boundary constants): 计算当前块在 y 维度上的最小和最大坐标 min_y, max_y。
2. 识别有效块 (Identify valid block ids):
* 在共享内存中定义一个 valid_block_ids 数组。
* 通过 gunc::reduce 并行归约，计算出当前块中所有线程 y 坐标的 block_min_y 和 block_max_y。
* 检查是否在边界内 (check if block is within y boundaries): 如果 block_min_y < min_y 且 block_max_y > max_y，则 valid_block_ids 记录当前 m_block_valid。
3. 异步预取 (async prefetch): 异步拷贝第一个有效块的数据。
4. 遍历所有有效块 (loop over all valid blocks):
* 对 valid_block_ids 中的每个有效块执行 compute() 和 apply_mask()。
* 在处理当前块的同时，异步预取下一个有效块的数据，实现计算和数据传输的重叠。

图示中的紫色点代表需要计算的块，而白色圆圈代表因完全在掩码区域之外而无需计算的块。这种方法可以提前剔除无效的计算块，从而提升效率。

Rab & dRab Cute 实践

Rab & dRab Cute 实践: Offset & Shape & Stride

本节介绍如何使用 Cute 库处理张量的形状、偏移和步长。

对于一个 torch 张量 [L, heads, dim]，可以进行布局变换。如图所示，一个逻辑上为长条形（2L * dim）的张量可以被重新排列为一个更宽的形状（L * 2dim），这在处理多头注意力中的批处理和数据重排时非常有用。

Rab & dRab Cute 实践: 获取子张量

如何获取你想要的子张量:
对于一个 [L, heads, dim] 的 torch 张量，有两种主要方法获取子张量：

1. 方法一 (切片和索引):
   • 偏移到当前 q-len 的起始点。
   • 构造一个 [q_len, head, dim] 的张量。
   • 使用 "local_tile" 来进行切片和索引。
2. 方法二 (偏移):
   • 另一个方法是偏移到当前 seq_len，继续偏移到当前 m_block，然后偏移到当前 head，所有东西都是通过偏移完成的。
3. 这两种方法可以组合使用。

代码与图示:
* 代码展示了如何使用 make_tensor 和 make_coord 来定义张量视图，通过指定形状和步长来访问内存的不同部分。
* 图示展示了内存的线性连续性，以及如何通过 head_stride 和 row_stride 来访问不同的头（h1, h2）和行。hx 表示属于第 x 个头的元素。

Rab & dRab Cute 实践: Rab-Offset & Shape & Stride

Stage 1 & Stage 2: 内存优化

• Stage 1:
  • rab 和 drab 的形状为 [bs, heads, max_seq_len_k, seqlen_k_rounded]。
  • 内存使用量为 bs * heads * max_k * max_k。
  • 由于 actual_seq_offset 的存在，导致了大量的冗余缓冲区（redundant_buffer），造成内存浪费。
• Stage 2:
  • rab 和 drab 的形状为 [bs, heads, max_seq_len_q, max_k]。
  • 内存使用量为 bs * heads * max_q * max_k。
  • 通过调整 shape 和 offset，显著减少了冗余缓冲区的大小，优化了内存使用。

关于填充 (Padding) 的说明:
进行填充的原因是 LDGSTS 和 TMA 执行 g2s (global to shared memory) 拷贝时，地址必须对齐到 16 字节。

Stage 3: 最终内存优化

• Stage 3:
  • rab 和 drab 的形状为 [heads, total_qk]，其中 total_qk 是所有批次和序列长度的总和。
  • 内存使用量为 heads * total_qk，完全没有冗余内存使用。
  • 这种布局还可以消除 dRab 的清零操作（通过使用 torch::empty 而非 torch::zeros），节省了延迟。
  • 代码中通过 offset_rab 和 actual_seq_k 来精确计算每个块在 total_qk 维度上的偏移。

Rab & dRab Cute 实践: dRab 计算与内存合并

使用 Cute 实现内存合并 (Memory Coalescing)

内存合并对于 GPU 性能至关重要。

• 左图 (Mem Coalescing Good!):

  • 写回方式: S->(R) G (从共享内存写回全局内存，可能经过寄存器)
  • 一个线程连续写入 8 个半精度浮点数 (half vals)。
  • 使用了 GmemTiledCopy 和 Copy_Atom<AutoVectorizingCopy<...>>，访问模式是连续的，实现了良好的内存合并，效率高。

• 右图 (Mem Coalescing Bad!):

  • 写回方式: R->G (从寄存器直接写回全局内存)
  • 一个线程连续写入 2 个半精度浮点数。
  • 使用了 GmemCopyAtom<dRab, ...>。由于线程写入的内存地址是不连续的（跨越多行），导致内存访问分散，无法合并，性能较差。图中 W0 和 W2 的访问模式清晰地展示了这种非合并访问。

使用 Cute 进行 dRab 计算与内存合并

本页幻灯片展示了在使用 Cute 库时，针对 dRab 计算的一种优化流程，即 Opt: R2S->S2R->R2G。此流程旨在通过数据布局的转换，最终实现高效的内存合并（Memory Coalescing）写入操作。

流程步骤如下：
1. 寄存器张量 (Reg Tensor, MMA 布局)：计算的初始数据位于寄存器中，采用为 MMA（Matrix Multiply-Accumulate）指令优化的布局。
2. 共享内存张量 (Smem Tensor, MMA 布局)：通过 R2S（Register to Shared Memory）操作，数据从寄存器转移到共享内存，布局保持为 MMA 布局。
3. 寄存器张量 (Reg Tensor, Tiledcopy 布局)：通过 S2R（Shared Memory to Register）操作，数据从共享内存加载回寄存器，但其布局被转换为 Tiledcopy 布局。
4. 全局内存张量 (Gmem Tensor, Tiledcopy 布局)：最后，通过 R2G（Register to Global Memory）操作，将采用 Tiledcopy 布局的寄存器数据写回全局内存。

此流程的关键在于，从 MMA 布局到 Tiledcopy 布局的转换，使得最终写入全局内存的数据是连续排列的，从而能够触发内存合并，显著提升写操作的带宽和效率。

使用 Cute 实现 AtomicAdd

本页幻灯片讨论了在 rab_head == 1 的情况下，如何使用 Cute 结合原子加（AtomicAdd）操作进行累加规约。根据累加精度的不同，提出了两种实现方案。

方案一：FP16 累加 (FP16 Accm)

此方案适用于头（heads）的数量较少的场景。

• 流程: S2G (Shared to Global Memory)。
• 实现要点:
  1. 当头数较少时，使用半精度（half precision）。
  2. 使用 __half2 或 __nv_bfloat162 内置函数来执行 atomicAdd，这允许一次原子操作处理两个半精度值。
  3. TiledCopy 操作可以被定义为一个线程对应两个半精度浮点值（two half vals）。

方案二：FP32 累加 (FP32 Accm)

此方案适用于头（heads）的数量较大的场景，以保证累加的精度。

• 流程: R2G -> G2R -> R2G。
• 实现要点:
  1. 当头数较大时，使用 fp32 精度。
  2. 使用标准的 float 类型来执行 atomicAdd。
  3. TiledCopy 操作可以被定义为一个线程对应一个单精度浮点值（one float val）。
• 优化提示: 图中特别指出，在此方案的 R2G 步骤中，同样可以应用内存合并优化（memory coalescing opt），以提升性能。

问答环节