1
import torch
2

3
def softmax(x, dim=-1):
4
    """
5
    softmax(z)_i = exp(z_i) / sum_j exp(z_j)
6
    """
7
    # 第 1 步：求最大值
8
    # 沿目标维度 dim 求最大值，keepdim=True 保持原维度数，
9
    # 便于后续广播减法。x_max 的形状会在 dim 维度变成 1，其他维不变。
10
    x_max = x.max(dim=dim, keepdim=True)[0]   # [0] 取出最大值
11

12
    # 第 2 步：平移输入数据，减去最大值
13
    # 数学上等价于将分子分母同除以 exp(x_max)，防止指数爆炸。
14
    # 因为 x_max 是沿 dim 维的最大值，x - x_max 的所有元素 ≤ 0，
15
    safe_x = x - x_max                     # 广播机制
16

17
    # 第 3 步：指数计算
18
    # 对平移后的张量逐元素求 exp。
19
    exp_x = torch.exp(safe_x)
20

21
    # 第 4 步：求和
22
    # 沿 dim 维度对 exp_x 求和，同样 keepdim=True 保持维度以便广播。
23
    sum_exp = exp_x.sum(dim=dim, keepdim=True)
24

25
    # 第 5 步：归一化得到最终概率
26
    output = exp_x / sum_exp
27

28
    return output
29

30
def naive_softmax(x, dim=-1):
31
    exp_x = torch.exp(x)
32
    return exp_x / exp_x.sum(dim=dim, keepdim=True)
33

34
if __name__ == "__main__":
35
    # 一维向量
36
    x1 = torch.tensor([1.0, 2.0, 3.0])
37
    p1 = softmax(x1)
38
    print("一维 softmax:\n",p1)
39
    print(f"Torch Softmax:\n{torch.softmax(x1, dim=-1)}")
40
    print("求和 =", p1.sum().item())  # 应为 1
41

42
    # 二维矩阵
43
    x2 = torch.tensor([[1., 2., 3.],
44
                       [4., 5., 6.]])
45
    p2 = softmax(x2, dim=-1)          # 对每一行做 softmax
46
    print("\n二维 softmax:\n",p2)
47
    print(f"Torch Softmax:\n{torch.softmax(x2, dim=-1)}")
48
    print("每行求和 =", p2.sum(dim=-1))  # 结果 [1, 1]

1
@triton.jit
2
def kernel_softmax_fuse(
3
    x_ptr, x_row_stride,
4
    y_ptr, y_row_stride,
5
    n_cols,
6
    BLOCK_SIZE: tl.constexpr,
7
):
8
    row_idx = tl.program_id(0)
9
    x_ptr += row_idx * x_row_stride
10
    y_ptr += row_idx * y_row_stride
11
    idx = tl.arange(0, BLOCK_SIZE)
12
    x = tl.load(x_ptr + idx, mask=idx < n_cols, other=-float("inf"))
13
    x = tl.exp(x - tl.max(x))
14
    eps = float(1e-9)
15
    x /= tl.maximum(tl.sum(x), eps)
16
    tl.store(y_ptr + idx, x, mask=idx < n_cols)
17

18

19
def triton_softmax_dim1_fuse(x):
20
    n_rows, n_cols = x.shape
21
    y = torch.empty_like(x)
22
    kernel_softmax_fuse[[n_rows]](
23
        x, x.stride(0),
24
        y, y.stride(0),
25
        n_cols,
26
        BLOCK_SIZE=triton.next_power_of_2(n_cols),
27
        num_warps=32,
28
    )
29
    return y

1
def stride(self, dim: None = None) -> tuple[_int, ...]:
2
        r"""
3
        stride(dim) -> tuple or int
4

5
        Returns the stride of :attr:`self` tensor.
6

7
        Stride is the jump necessary to go from one element to the next one in the
8
        specified dimension :attr:`dim`. A tuple of all strides is returned when no
9
        argument is passed in. Otherwise, an integer value is returned as the stride in
10
        the particular dimension :attr:`dim`.
11

12
        Args:
13
            dim (int, optional): the desired dimension in which stride is required
14

15
        Example::
16

17
            >>> x = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
18
            >>> x.stride()
19
            (5, 1)
20
            >>> x.stride(0)
21
            5
22
            >>> x.stride(-1)
23
            1
24
        """

1
def triton_softmax_dim1_fuse(x):
2
    n_rows, n_cols = x.shape
3
    y = torch.empty_like(x)
4
    kernel_softmax_fuse[[n_rows]](
5
        x,
6
        x.stride(0),
7
        y,
8
        y.stride(0),
9
        n_cols,
10
        BLOCK_SIZE=triton.next_power_of_2(n_cols),
11
        num_warps=32,
12
    )
13
    return y
14

15

16
@triton.jit
17
def kernel_softmax_tile(
18
    x_ptr, x_row_stride,
19
    y_ptr, y_row_stride,
20
    n_cols,
21
    BLOCK_SIZE: tl.constexpr,
22
    CACHE_OPT: tl.constexpr,
23
):
24
    row_idx = tl.program_id(0)
25
    x_ptr += row_idx * x_row_stride
26
    y_ptr += row_idx * y_row_stride
27

28
    mm = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - float("inf")
29

30
    for i in range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
31
        idx = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
32
        x = tl.load(x_ptr + idx, mask=idx < n_cols, other=-float("inf"))
33
        mm = tl.maximum(mm, x)
34
    mm = tl.max(mm)
35

36
    ss = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
37

38
    if CACHE_OPT:
39
        for i in range(tl.cdiv(n_cols, BLOCK_SIZE) - 1, -1, -1):
40
            idx = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
41
            x = tl.load(x_ptr + idx, mask=idx < n_cols, other=-float("inf"))
42
            x = tl.exp(x - mm)
43
            ss += x
44
    else:
45
        for i in range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
46
            idx = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
47
            x = tl.load(x_ptr + idx, mask=idx < n_cols, other=-float("inf"))
48
            x = tl.exp(x - mm)
49
            ss += x
50

51
    ss = tl.sum(ss)
52
    eps = float(1e-9)
53
    ss = tl.maximum(ss, eps)
54

55
    for i in range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
56
        idx = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
57
        x = tl.load(x_ptr + idx, mask=idx < n_cols, other=-float("inf"))
58
        x = tl.exp(x - mm) / ss
59
        tl.store(y_ptr + idx, x, mask=idx < n_cols)
60

61

62
def triton_softmax_dim1_tile(x, cache_opt=True):
63
    n_rows, n_cols = x.shape
64
    y = torch.empty_like(x)
65
    kernel_softmax_tile[[n_rows]](
66
        x, x.stride(0),
67
        y, y.stride(0),
68
        n_cols,
69
        BLOCK_SIZE=2**14,
70
        CACHE_OPT=cache_opt,
71
        num_warps=32,
72
    )
73
    return y