File size: 2,382 Bytes
0d27699 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | # -*- coding: utf-8 -*-
# Copyright (c) 2023-2024, Songlin Yang, Yu Zhang
from typing import Optional
import torch
import triton
import triton.language as tl
from fla.ops.utils.op import exp
@triton.autotune(
configs=[
triton.Config({}, num_warps=1),
triton.Config({}, num_warps=2),
triton.Config({}, num_warps=4),
triton.Config({}, num_warps=8),
triton.Config({}, num_warps=16),
triton.Config({}, num_warps=32)
],
key=['D']
)
@triton.jit
def softmax_fwd_kernel(
x,
p,
D: tl.constexpr,
B: tl.constexpr
):
i_n = tl.program_id(0)
o_d = tl.arange(0, B)
m_d = o_d < D
b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
b_m = tl.max(b_x, 0)
b_x = exp(b_x - b_m)
b_p = b_x / tl.sum(b_x, 0)
tl.store(p + i_n * D + o_d, b_p.to(p.dtype.element_ty), mask=m_d)
@triton.autotune(
configs=[
triton.Config({}, num_warps=1),
triton.Config({}, num_warps=2),
triton.Config({}, num_warps=4),
triton.Config({}, num_warps=8),
triton.Config({}, num_warps=16),
triton.Config({}, num_warps=32)
],
key=['D']
)
@triton.jit
def softmax_bwd_kernel(
p,
dp,
ds,
D: tl.constexpr,
B: tl.constexpr
):
i_n = tl.program_id(0)
o_d = tl.arange(0, B)
m_d = o_d < D
b_p = tl.load(p + i_n * D + o_d, mask=m_d, other=0.)
b_dp = tl.load(dp + i_n * D + o_d, mask=m_d, other=0.)
b_pp = tl.sum(b_p * b_dp, 0)
b_ds = b_p * b_dp - b_p * b_pp
tl.store(ds + i_n * D + o_d, b_ds.to(ds.dtype.element_ty), mask=m_d)
def softmax_fwd(
x: torch.Tensor,
dtype: Optional[torch.dtype] = torch.float
) -> torch.Tensor:
shape = x.shape
x = x.view(-1, x.shape[-1])
N, D = x.shape
B = triton.next_power_of_2(D)
p = torch.empty_like(x, dtype=dtype)
softmax_fwd_kernel[(N,)](
x=x,
p=p,
D=D,
B=B
)
return p.view(*shape)
def softmax_bwd(
p: torch.Tensor,
dp: torch.Tensor,
dtype: Optional[torch.dtype] = torch.float
) -> torch.Tensor:
shape = p.shape
p = p.view(-1, p.shape[-1])
ds = torch.empty_like(p, dtype=dtype)
N, D = p.shape
B = triton.next_power_of_2(D)
softmax_bwd_kernel[(N,)](
p=p,
dp=dp,
ds=ds,
D=D,
B=B
)
return ds.view(*shape)
|