ops(6)：embedding 层与 LM head 层的 CUDA 实现

原文链接: https://zhuanlan.zhihu.com/p/695785781
发布时间: 2024-05-11
作者: ifromeast

概述

本文详细介绍了 Transformer 模型中 embedding 层和 LM head 层的 CUDA 实现与优化。这两个层分别构成了整个模型的首和尾，在模型性能中扮演着重要角色。

一、Embedding 层的 CUDA 实现

1.1 Embedding 与 Linear 的区别

基本概念

在语言模型中，embedding 层的作用是将文本的离散语义和位置信息转换为高维向量表示。在标准 GPT 模型中：

• word embedding (wte): 词嵌入矩阵
• position embedding (wpe): 位置嵌入矩阵

Embedding 的本质

Embedding 操作本质上是向量索引过程，可以等价转换为矩阵乘法：

索引方式：

import torch
idx = torch.tensor([2, 3, 1])
embedding = torch.nn.Embedding(4, 5)
output = embedding(idx)  # 直接索引

矩阵乘法方式：

# 将索引转换为 one-hot 编码
onehot = torch.nn.functional.one_hot(idx)
# 使用线性层（权重为 embedding 矩阵的转置）
linear = torch.nn.Linear(4, 5)
linear.weight = torch.nn.Parameter(embedding.weight.T.detach())
output = linear(onehot.float())  # 矩阵乘法

数学表达式：

z = x * W = [W_s1 ... W_sD]

其中 x 是 one-hot 编码，W 是 embedding 矩阵。

反向传播推导

梯度计算：

∂J/∂W_ij = Σ_k (∂J/∂z_k) * (∂z_k/∂W_ij)

关键性质：

• 当 k ≠ j 时，∂z_k/∂W_ij = 0
• 当 k = j 时，∂z_k/∂W_ij = 1

因此：

∂J/∂W_sj = ∂J/∂z_j

写成矩阵形式，反向传播是对梯度的索引过程。

1.2 Embedding 层的实现

CPU 基准实现

void encoder_forward_cpu(float* out,
                   const int* inp, const float* wte, const float* wpe,
                   int B, int T, int C) {
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            float* out_bt = out + b * T * C + t * C;
            int ix = inp[b * T + t];
            const float* wte_ix = wte + ix * C;
            const float* wpe_t = wpe + t * C;
            for (int i = 0; i < C; i++) {
                out_bt[i] = wte_ix[i] + wpe_t[i];
            }
        }
    }
}

实现逻辑：

1. 遍历 batch (B) 和 sequence (T) 维度
2. 根据输入索引 ix 获取对应的 word embedding
3. 根据位置 t 获取对应的 position embedding
4. 将两者相加得到最终输出

CUDA 实现 V1：基础并行化

优化策略：

• 在 B×T 维度上并行
• 每个线程处理一个 (b, t) 位置的所有 C 个元素

性能数据：

block_size   32 | time 0.3918 ms | bandwidth 256.93 GB/s
block_size   64 | time 0.3921 ms | bandwidth 256.73 GB/s
block_size  128 | time 0.4000 ms | bandwidth 251.64 GB/s
block_size  256 | time 0.7029 ms | bandwidth 143.21 GB/s

CUDA 实现 V2：完全并行化

优化策略：

• 在 B×T×C 三个维度上完全并行
• 每个线程只处理一个元素
• 索引操作独立，适合高度并行

性能数据：

block_size   32 | time 0.3366 ms | bandwidth 299.10 GB/s
block_size   64 | time 0.1600 ms | bandwidth 629.16 GB/s
block_size  128 | time 0.0853 ms | bandwidth 1180.48 GB/s
block_size  256 | time 0.0745 ms | bandwidth 1350.73 GB/s
block_size  512 | time 0.0746 ms | bandwidth 1350.01 GB/s

性能提升：相比 V1，带宽提升约 5 倍。

CUDA 实现 V3：向量化内存访问

优化策略：

• 使用 128 位向量化加载/存储（x128）
• 一次处理多个元素，减少内存事务
• 使用 load128cs 和 store128 优化内存访问模式

性能数据：

block_size   32 | time 0.0573 ms | bandwidth 1756.26 GB/s
block_size   64 | time 0.0538 ms | bandwidth 1869.95 GB/s
block_size  128 | time 0.0537 ms | bandwidth 1874.82 GB/s
block_size  256 | time 0.0536 ms | bandwidth 1879.17 GB/s
block_size  512 | time 0.0539 ms | bandwidth 1868.99 GB/s

性能提升：相比 V2，带宽再提升约 40%，达到 1879 GB/s。

1.3 Embedding 反向传播实现

CUDA 实现：使用原子操作

关键技术：

• 使用 atomicAdd 避免多个线程同时写入同一位置造成的竞争条件
• 多个不同位置可能索引到同一个 embedding 向量

性能数据：

block_size   32 | time 0.3296 ms
block_size   64 | time 0.1735 ms
block_size  128 | time 0.1719 ms
block_size  256 | time 0.1728 ms

二、LM Head 层的实现

2.1 LM Head 的作用

在 Transformer 的 Causal Language Model 中，LM head 层负责：

1. 推理阶段：

   • 根据 logits 通过 softmax 计算概率分布 probs
   • 使用解码算法（如 greedy、beam search、sampling）生成 next token

2. 训练阶段：

   • 根据 probs 和 labels 计算 loss
   • 计算 logits 的梯度 dlogits，开始反向传播

2.2 融合的 Classifier Kernel

关键技术点

1. Softmax 计算

• 使用 warp 级别的协作组进行高效规约
• 计算 softmax 的 offset 和 scale 参数

2. Loss 计算

• 交叉熵损失：-log(prob)
• 只需要计算目标位置的概率

3. 梯度计算

• Softmax + CrossEntropy 的梯度：(prob - indicator) * dloss
• indicator 在目标位置为 1，其他位置为 0

4. 向量化内存访问

• 使用 float4 一次读取 4 个元素
• __ldcs 指示编译器使用流式缓存提示

5. 融合操作

• 在一个 kernel 中完成 softmax、loss 计算和梯度计算
• 减少内存访问次数，提高缓存利用率
• 可选地输出 probs（用于推理或调试）

性能数据

block_size   32 | time 10.460636 ms
block_size   64 | time 10.547235 ms
block_size  128 | time 9.895740 ms
block_size  256 | time 9.381703 ms
block_size  512 | time 9.059260 ms
block_size 1024 | time 8.828804 ms

性能特点：

• 该 kernel 是 DRAM 限制的（memory-bound）
• 较大的 block_size 性能更好
• 融合操作显著减少了内存带宽需求

三、优化技术总结

3.1 并行化策略

关键洞察：

• Embedding 操作的索引是独立的，适合高度并行
• 增加并行度可以显著提高性能
• 向量化内存访问进一步提升带宽利用率

3.2 内存访问优化

1. 向量化加载/存储

• 一次处理 128 位（多个元素）
• 减少内存事务数量
• 提高内存带宽利用率

2. 缓存提示

• __ldcs：流式加载，数据不会长期保留在缓存中
• 适用于只使用一次的数据
• 为其他数据腾出缓存空间

3. 合并访问

• 确保相邻线程访问连续内存
• 利用 GPU 的内存合并机制
• 减少内存事务数量

3.3 原子操作

使用场景：

• Embedding 反向传播中的梯度累加
• 多个位置可能索引到同一个 embedding 向量

性能影响：

• 原子操作有性能开销
• 但对于 embedding 反向传播是必需的
• 冲突较少时性能影响可接受

3.4 Kernel 融合

LM Head 融合操作：

1. Softmax 计算
2. Loss 计算
3. 梯度计算

优势：

• 减少 kernel 启动开销
• 减少中间结果的内存读写
• 提高数据局部性和缓存利用率
• 整体性能提升显著

四、性能分析

4.1 Embedding 层性能演进

分析：

• V1→V2：并行度提升是主要因素
• V2→V3：内存访问模式优化
• V3 接近硬件理论带宽上限

4.2 Block Size 选择

Embedding 层：

• 最佳 block_size：128-512
• 过小：并行度不足
• 过大：寄存器/共享内存压力

LM Head 层：

• 最佳 block_size：512-1024
• 较大的 block_size 更好
• 因为是 memory-bound 操作

4.3 内存带宽利用率

理论分析：

• Embedding 是典型的 memory-bound 操作
• 计算量小，主要瓶颈在内存访问
• 优化目标：最大化内存带宽利用率

实际表现：

• V3 版本达到 1879 GB/s
• 接近 GPU 理论带宽（取决于具体硬件）
• 进一步优化空间有限

五、实现要点

5.1 Embedding 层

前向传播：

1. 根据输入索引获取 word embedding
2. 根据位置获取 position embedding
3. 两者相加得到输出
4. 使用向量化内存访问优化

反向传播：

1. 将输出梯度累加到对应 embedding 位置
2. 必须使用原子操作避免竞争
3. 这是索引操作的反向过程

5.2 LM Head 层

融合 Kernel 设计：

1. 使用 warp 级别协作组计算 softmax
2. 单线程计算 loss（避免同步）
3. 并行计算所有位置的梯度
4. 可选输出 probs（推理/调试）

数值稳定性：

• Softmax 使用 offset 技巧避免溢出
• Loss 计算使用 log-sum-exp 技巧

5.3 代码组织

参数说明：

• B: batch size
• T: sequence length
• C: embedding dimension
• V: vocabulary size
• P: padded vocabulary size（用于对齐）

内存布局：

• 使用行优先（row-major）布局
• 确保内存访问的连续性
• 便于向量化操作

六、参考资源

代码仓库：

• encoder_forward.cu
• encoder_backward.cu
• classifier_fused.cu

参考项目：

• karpathy/llm.c - 主要参考实现
• LLMs-from-scratch - Embedding 原理讲解

七、关键收获

1. Embedding 本质：索引操作，可等价为矩阵乘法，但索引更高效

2. 并行化策略：根据操作特点选择合适的并行维度，Embedding 适合完全并行

3. 内存优化：向量化访问、缓存提示、合并访问是提升性能的关键

4. Kernel 融合：减少内存访问，提高数据局部性，显著提升性能

5. 原子操作：在必要时使用，权衡性能和正确性

6. 性能分析：理解操作是 compute-bound 还是 memory-bound，针对性优化

总结：本文通过详细的实现和优化过程，展示了如何从基础版本逐步优化到高性能实现，最终达到接近硬件理论性能的水平。这些技术和思路可以推广到其他 CUDA 算子的实现中。