Softmax 稳定形式公式#

对于一行输入：

Softmax 的定义是：

可以拆分成两步：

• 对每一个元素取指得到正数
• 除以这些正数的总和，使得输出之和为 1

对于计算机而言，我们通常会先求输入的 max，因为计算机能表示的数值范围有限，如果输入元素有 1000，expf(100) 会溢出，所以解决办法是：所有元素先减去同一个最大值，再取指

即 Softmax 的稳定形式为：

分层归约思路#

在 CUDA 中，我们通常采用分层归约进行 Softmax

先考虑一个常见方案：一个 block 处理一整行输入。

假设一行有 1024 个元素，一个 block 256 个线程，每一个线程处理 4 个元素，即

1. 第一步，每一个线程先求自己的局部最大值：

其中，$I_t$ 表示线程 t 负责的下标集合，然后把所有吸纳成的局部最大值进行最大值归约，得到本行最大值 $m$：

2. 每一个线程都用同一个 m，计算自己的局部指数和，再归约所有的局部和：

3. 最后一步，每一个线程计算自己负责的输出

如果一个把一行分成若干块，每一个块求出自己的局部最大值与局部和，最终需要归约为整体最大值与局部和

需要注意的是：不同块的局部和不能直接相加，需要先把他们换成同一个基准：

因此，正确的整体指数和：

缩放点积注意力公式#

这是 Transformer 中使用的一种注意力计算，下面分别讲解一下每一个变量的含义：

其中：$K$ 的第 $j$ 行与 $V$ 的第 $j$ 行是一对，用 Key 判断匹配程度，用 Value 提供最终参与加权的内容。注意力的输出就是这些 Value 的加权和

1. 计算匹配分数：$S=QK^T/\sqrt{d}$#

取第 $i$ 个 Query（Q 的第 $i$ 行），与第 $j$ 个 Key（$K^T$ 的第 $j$ 列） 做点积：

$S_{ij}$ 越大，表示按照这里的点积评分方式，Query $i$ 与 key $j$ 越匹配。一行 $S_{i,:}$ 就是一个 Query 对全部 $N$ 个 Key 的评分

除以 $\sqrt d$ 是为了控制分数的尺度：在常见的统计假设下，点积的方差随 $d$ 增大，缩放可以避免 Softmax 过于尖锐。

NOTE

为什么随着 $d$ 增大，方差变大？ 对于点积公式，假设各分量相互独立，均值为 0，方差为 1，相互独立的随机变量相加，方差也相加，因此 $\text{Var}(s) = \Sigma\text{Var}(q_rk_r)=d$ 所以方差随着 $d$ 增长，标准差随着 $\sqrt{d}$ 增长，标准差更直接地反映结果的典型波动幅度。 因此除以 $\sqrt{d}$ 能让不同维度下的评分保持相近的波动尺度，避免维度增大时评分差距变得过大，使 Softmax 权重过度集中。

2. 对每行做 Softmax：将评分变成权重#

固定 Query $i$，对他的 $N$ 个评分进行稳定 Softmax：

得到的权重 $P$ 仍然为 $M \times N$，并且每一行：

这里要分清两个维度：

• 点积沿 $d$ 维求和：计算一对 Query、Key 的分数。
• Softmax 沿 $N$ 维归一化：比较同一个 Query 对所有 Key 的分数。

3. 对 Value 加权求和：$O=PV$#

对于 Query $i$，用刚才得到的权重，对所有 Value 向量进行加权:

也就是说，一个权重 $P_{ij}$ 乘的是整个 Value 向量 $V_{j,:}$，其中所有的分量使用同一个权重。

背景#

• Q 矩阵大小：$M \times d$，K 与 V 矩阵大小：$N \times d$
• $1 \leq$ M，N $\leq 100,000$
• $1 \leq$ d $\leq 128$
• 这里先采用 M $= 512$，N $= 256$

第一个版本先约定：

1
constexpr int BLOCK_SIZE = 256;
2
constexpr int NUM_WARPS = BLOCK_SIZE / 32;

分层归约#

在 Softmax 中，我们需要对每一个 block 求出当前的局部和与局部 max，需要用到分层归约的技巧

首先通过 __shfl_down_sync 实现 warp 内线程的互相通信

1
template <bool IS_MAX>
2
__device__ __forceinline__ float warp_reduce(float value)
3
{
4
    for (int offset = 16; offset > 0; offset >>= 1)
5
    {
6
        const float other = __shfl_down_sync(0xffffffffu, value, offset);
7
        if (IS_MAX)
8
            value = fmaxf(value, other);
9
        else
10
            value += other;
11
    }
12
    return value;
13
}

block_reduce 调用 warp_reduce，实现一个 Block 归约的输出

1
template <bool IS_MAX>
2
__device__ __forceinline__ float block_reduce(float value)
3
{
4
    __shared__ float shared[NUM_WARPS];
5
    const int warp_id = threadIdx.x / 32;
6
    const int lane = threadIdx.x % 32;
7

8
    value = warp_reduce<IS_MAX>(value);
9
    if (lane == 0)
10
        shared[warp_id] = value;
11

12
    __syncthreads();
13

14
    if (warp_id == 0)
15
    {
16
        const float identity = (IS_MAX) ? -CUDART_INF_F : 0.0f;
17
        value = (lane < NUM_WARPS) ? shared[lane] : identity;
18
        value = warp_reduce<IS_MAX>(value);
19

20
        if (lane == 0)
21
            shared[0] = value;
22
    }
23
    __syncthreads();
24

25
    value = shared[0];
26

27
    __syncthreads();
28
    return value;
29
}

核心 Kernel#