LLM Systems Assignment 4: Transformer CUDA Acceleration

本次 homework 主要内容是：实现 softmax 与 LayerNorm kernel fusion 形式的 fwd 与 bwd，完成 CUDA kernel 编写；接入 Python 层；接入 Transformer 模型并验证训练加速。

下面对本次 homework 实现时一些注意点进行介绍。

1. Softmax

1.1 forward

课程已经提供了 block_dim <= 32 时的 kernel 实现代码，我们只需要实现支持 block_dim > 32 版本的 kernel。

在开始实现 kernel 之前，我们首先要理解 kernel 的 thread/block 组织以及它们如何协作完成 softmax forward 计算。

在 launch_attn_softmax 这个 C 语言接口函数中进行了 kernel launch，查看其中的代码可以理解线程任务划分：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
void launch_attn_softmax(float* inp, const float* attn_mask, int batch_size, int nhead, int from_len, int to_len,
                         bool mask_future, cudaStream_t stream) {
  constexpr int float_size = sizeof(float);
  const int inp_size = batch_size * nhead * from_len * to_len * float_size;
  // Ignore the nhead & from_len dim of attn_mask since the mask is the same among heads/queries.
  const int attn_mask_size = batch_size * to_len * float_size;

  float *d_inp = nullptr, *d_attn_mask = nullptr;

  CUDA_CALL(cudaMalloc((void**)&d_inp, inp_size));
  CUDA_CALL(cudaMemcpy(d_inp, inp, inp_size, cudaMemcpyHostToDevice));

  if (attn_mask != nullptr) {
    CUDA_CALL(cudaMalloc((void**)&d_attn_mask, attn_mask_size));
    CUDA_CALL(cudaMemcpy(d_attn_mask, attn_mask, attn_mask_size, cudaMemcpyHostToDevice));
  }

  // The attention score tensor shape is [batch_size, nhead, from_len, to_len], the former `batch_size` and `nhead` dim
  // is independent(no need cooperate between threads in a block), so make them coresponding to the `gridDim.y` and
  // `gridDim.z` when launch kernel. The `from_len` dim is also independent, but it's size usually related to
  // `to_len`(i.e., `from_len` <= `to_len`), so we set `gridDim.x` to a value less than `to_len` when dispatch and
  // launch kernel. For each row in `from_len` dim, we need to reduce among the `to_len` elements, the threads work of
  // block is:
  //    - Each thread read contiguous 1 or 2 (the third template parameter) values.
  //    - Each thread loop over the `from_len` dim to handle multiple queries.
  // `nhead` is `gridDim.z`, which means different block (blockIdx.z) computes different head_id.
  // `batch_size` is `gridDim.y`, which means different block (blockIdx.y) computes different batch_id.
  dim3 grid_dim(1, batch_size, nhead);
  if (to_len <= 32) {
    // blockDim = 32, only 1 warp in each block. Each warp handles the full
    ker_attn_softmax_lt32<float, 32, 1><<<grid_dim, 32, 0, stream>>>(d_inp, d_attn_mask, from_len, to_len, mask_future);
  } else if (to_len <= 64) {
    ker_attn_softmax_lt32<float, 32, 2><<<grid_dim, 32, 0, stream>>>(d_inp, d_attn_mask, from_len, to_len, mask_future);
  } else if (to_len <= 128) {
    // Increase gridDim.x to 16 (each block handles at most 128/16 = 8 rows), blockDim.x = 64, 2 warps in each block.
    grid_dim.x = 16;
    ker_attn_softmax<float, 64, 2><<<grid_dim, 64, 0, stream>>>(d_inp, d_attn_mask, from_len, to_len, mask_future);
  } else if (to_len <= 256) {
    grid_dim.x = 32;
    ker_attn_softmax<float, 128, 2><<<grid_dim, 128, 0, stream>>>(d_inp, d_attn_mask, from_len, to_len, mask_future);
  } else if (to_len <= 512) {
    grid_dim.x = 64;
    ker_attn_softmax<float, 256, 2><<<grid_dim, 256, 0, stream>>>(d_inp, d_attn_mask, from_len, to_len, mask_future);
  } else if (to_len <= 1024) {
    grid_dim.x = 128;
    ker_attn_softmax<float, 512, 2><<<grid_dim, 512, 0, stream>>>(d_inp, d_attn_mask, from_len, to_len, mask_future);
  } else {
    throw std::runtime_error("Sequence length greater than 512 is currently not supported");
  }

  // 略 ...
}

kernel 的线程任务划分已经在上述代码注释中进行了说明，这里重复解释如下：

  • 输入 tensor 形状为 (batch_size, nhead, from_len, to_len)，其中 batch_size 与 nhead 都很好理解，就是训练时 batch_size 与 Attention 的 head 数量。from_len 与 to_len 则分别表示 query 与 key 的 token 数量。由于 minitorch 是一个仅考虑训练的教学性质框架，实际上在我们的 homework 中进行 GPT-2 模型训练时，softmax kernel 的输入参数一定满足 from_len == to_len。
  • 在理解输入 tensor 形状的基础上，softmax forward kernel 要做的事情自然很清楚了：沿着 (batch_size, nhead, from_len) 三个互相独立的维度，在 to_len 维度做 to_len 个元素的 reduce，进行 softmax 计算。
  • 然后我们就可以映射输入 tensor shape 到 launch kernel 时的 blockDim/gridDim 参数了：
    • gridDim 的 y 与 z 取值分别为 batch_size 与 nhead，表示输入 tensor 最外层独立的两个维度；
    • gridDim 的 x 取值与 from_len 有关。虽然 from_len 维度与 softmax 计算过程也是独立的，但是由于 from_len 通常小于等于 to_len（也就是 query token 数量肯定小于等于 key token 数量），因此在实际进行 kernel launch 时，gridDim.x 的取值根据 to_len 的大小进行调整。这样，相当于一个 block 内的线程可能要负责计算同一个 (batch_id, head_id)（对应 blockIdx.y, blockIdx.z）内部的多个 from_len 维度的行。
    • blockDim.x 的取值取决于 to_len，即 reduce 维度的大小。以上述代码中 to_len <= 128 的 if 分支为例，launch kernel 时设置了 blockDim.x = 64，这意味着一个 block 内的 64 个线程需要协作最多处理 to_dim 的 128 个元素，也就是一个线程最多处理 128 / 64 = 2 个元素。

在理解 kernel 线程组织方式的基础上，实现 softmax forward kernel 的流程就很清晰了：

  • 加载当前 block 负责的位置对应的 attn_mask (长度为 to_len)到寄存器；
  • 计算当前 block/thread 负责的输入地址；
  • 线程循环处理负责的不同行；
    • 读取当前处理行到寄存器；
    • 计算 local max 然后 reduce 得到 block max（需要同时考虑 attn_mask）；
    • 计算 exp(val - max_val) 然后求和得到 softmax 分母（sum(exp(val - max_val))）；
    • 计算最终的 softmax 结果，写到结果 global memory；

我们需要实现的 kernel 与已经提供的 kernel 整体流程基本一致，唯一区别是进行 block 内 reduce 时，是只需要 warp 级别同步；还是需要 block 级别同步（借助 shared memory 得到 reduce 结果）。

1.2 backward

softmax 的 backward kernel 代码已经提供，这里简要介绍下 softmax backward kernel 整体流程。

首先看 kernel launch 的 blockDim/gridDim 参数：

1
2
3
const int warps_per_block = 4;
dim3 grid_dim((rows + warps_per_block - 1) / warps_per_block);
dim3 block_dim(WARP_SIZE, warps_per_block);

可以看到，backward kernel 的 thread/block 组织方式简单很多：

  • 一个 block 固定是 32 * 4 共 128 线程。
    • 每个 warp 负责处理一行；一个 block 内 4 个 warp 处理 4 行。
    • 一个 warp 内的 32 线程负责读取并计算 to_len （函数参数的 softmax_len）个数值。
  • gridDim.x 数量由总 rows = batch_size * nhead * from_len 除以每个 block 处理 4 行向上取整计算得到。
  • 在我们需要实现的根据 softmax_len 大小进行 kernel launch 时，我们要计算每个线程需要读取的最大元素数量填充 backward kernel 的 ITERATIONS 模板参数。

kernel 内的执行流程为：

  • 计算当前 warp 负责的 row_id；
  • 每个线程循环读取上层梯度和原始输入到寄存器，并进行线程局部 reduce；
  • 进行 warp 级别 reduce；
  • 根据 softmax backward 公式计算最终梯度并写到 global memory。

2. LayerNorm

2.1 forward

LayerNorm forward kernel 的流程相对而言比较直观：

  • 根据课程 slides 的方法一次遍历 + reduce 同时得到 mean/var；
  • 利用 shared memory 广播 mean/var 到每个 thread（这一步不可省略，因为已有代码提供的 blockReduce 结果只能让 block 的第一个 warp 得到正确 reduce 结果）；
  • 再次遍历，根据 LayerNorm 公式计算结果并写到 global memory。

有几点要注意的是：

  • blockReduce 若连续多次调用，中间必须加 block 级别同步保证前一次的 reduce 完成，否则可能出现竞态得到错误结果。
  • kernel 参数传入的 hidden_dim 是真实 dim 除以 4 之后的，以方便 kernel 内部指针偏移量计算；但是在计算 mean/var 时，要除以真正的 dim。

2.2 backward

LayerNorm 的 backward 实现需要计算两个梯度：

  • 输入的 (rows, hidden_dim) shape tensor 的梯度。
  • LayerNorm 计算中的 weight 与 bias 的梯度，也就是下面的 dgamma_dbeta kernel 完成的功能。

2.2.1 d_inp

这一 kernel 主要实现流程就是按照课程 slides 或者代码注释完成：

  • 每个线程循环遍历得到 LayerNorm backward 公式中的两个 local reduce 结果；
  • 调用 blockReduce 并广播得到整个 block 的 reduce 结果；
  • 再次循环遍历，计算 backward 结果。由于向量访存优化的存在，实现时可能会有些重复代码，可以考虑抽取为单独的 device 函数以简化代码，例如：
1
2
3
4
5
6
7
__device__ __forceinline__ float ln_grad(const float& x, const float& dy, const float& w, const float& rstd,
                                         const float& std, const float& mean, const float& sum1, const float& sum2,
                                         const int& m) {
  const float rmstd = 1.0f / (m * std);
  const float xhat = (x - mean) * rstd;
  return (w * dy) * rstd - rmstd * (sum1 + xhat * sum2);
}

2.2.2 dgamma_dbeta

homework 文档已经提供了 dgamma_dbeta 的计算公式，我们只需要按照公式实现。在看 kernel 具体流程之前，我们先看下这个 kernel 的线程任务划分。

首先是 launch kernel 代码：

1
2
3
4
5
// This calculates the number of blocks needed to cover the data along the specified dimension, rounds it up.
dim3 grid_dim((hidden_dim + TILE_DIM - 1) / TILE_DIM);
dim3 block_dim(TILE_DIM, TILE_DIM);
ker_ln_bw_dgamma_dbetta<float><<<grid_dim, block_dim, 0, stream_1>>>(
    d_gamma_grad, d_betta_grad, d_out_grad, d_inp, d_gamma, d_betta, d_vars, d_means, batch_size, hidden_dim);

其中 blockDim 为固定的 32 * 32；gridDim 则用 hidden_dim 对 32 向上取整得到。由此我们可以看出来线程任务划分如下：

  • 不同 block 处理整个 [rows * hidden_dim] 输入的 [:, blockIdx.x * TILE_DIM: blockIdx.x * TILE_DIM + TILE_DIM 部分，也就是除了整个输入所有行的部分列；
  • 对一个 block 来说，线程划分为了 x 和 y 两个维度，其中：
    • 对于 threadIdx.y 相同的一个 warp 内 32 个线程，协作读取输入的一行；
    • threadIdx.y 不同的 warp 之间负责不同 rows；

kernel 整体流程如下：

  • 线程跨步读取不同 rows，每个 thread 计算得到负责的 column_id 的多个 rows 线程局部 reduce 结果（注意边界判断）；这一步完成后，不同 warp 得到了不同 rows subset 的 reduce 结果，warp 内的一个线程拥有某个 column_id 的部分 rows reduce 结果；
  • 使用 shared memory 写入 local reduce 结果并转置读取；这样，一个 warp 的 32 个线程拥有了同一个 column_id 的不同 rows subset reduce 结果；
  • 然后利用 cooperative_groups 接口进行 warp 内 reduce，得到了某个 column_id 的全部 rows 最终 reduce 结果；
  • 最后，让不同 warp 内的 lane_id == 0 的第一个线程，写最终 reduce 结果到结果 global memory。

3. Adopt Fused Kernels in Transformer

这一步的整体流程如下：

  • 给相关 Module 构造函数增加 use_fused_kernel 参数；
  • 修改相关 Module forward 实现，在 use_fused_kernel 开启时，调用 tensor_functions 模块的 xxxFunction.apply 接口，实现 fused kernel 调用。

需要注意的是，这一步可能要从之前的 homework3 copy 部分代码，需要仔细 check 下 copy 过来的代码是完整的。

最后，附上是否开启 kernel fuse 的运行耗时记录：

1
2
3
4
5
6
7
8
9
10
11
12
# 不开启 kernel fuse

Training (epoch 0 / 1):  99%|█████████▉| 156/157 [30:15<00:10, 10.81s/it, loss=5.77, lr=0.0002, tokens_per_sec=318]
Training (epoch 0 / 1): 100%|██████████| 157/157 [30:15<00:00,  8.75s/it, loss=5.77, lr=0.0002, tokens_per_sec=318]
Training (epoch 0 / 1): 100%|██████████| 157/157 [30:15<00:00, 11.56s/it, loss=5.77, lr=0.0002, tokens_per_sec=318]

# 开启 kernel fuse

Training (epoch 0 / 1):  99%|█████████▉| 156/157 [26:47<00:09,  9.82s/it, loss=5.8, lr=0.0002, tokens_per_sec=364]
Training (epoch 0 / 1): 100%|██████████| 157/157 [26:47<00:00,  7.91s/it, loss=5.8, lr=0.0002, tokens_per_sec=364]
Training (epoch 0 / 1): 100%|██████████| 157/157 [26:47<00:00, 10.24s/it, loss=5.8, lr=0.0002, tokens_per_sec=364]

由于 softmax/LayerNorm 仅占 Transformer 计算的一小部分，所以整体性能提升有限，但是整体符合预期。

4. Conclusion

自定义 kernel fuse 流程：

  • 实现 CUDA kernel；
  • 实现 C++ Operator 接口；
  • 接入 Python 保留接口到上层；
  • 自定义 Function 实现 forward/backward 方法；
  • 实现 Module forward 方法，内部调用自定义 Function apply 方法，实现 fused kernel 调用。
ML System
#CMU 11868 #LLM System #kernel
LLM Systems Assignment 4: Transformer CUDA Acceleration
https://arcsin2.cloud/posts/2026/10/1691777268/
作者
arcsin2
发布于
2026年10月17日
许可协议