File size: 1,903 Bytes
8cbec64 | 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 | # -*- 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, log
@triton.heuristics({
'HAS_SCALE': lambda args: args['scale'] is not None
})
@triton.autotune(
configs=[
triton.Config({}, num_warps=num_warps)
for num_warps in [1, 2, 4, 8, 16, 32]
],
key=['D']
)
@triton.jit
def logsumexp_fwd_kernel(
x,
z,
scale,
D: tl.constexpr,
B: tl.constexpr,
HAS_SCALE: tl.constexpr
):
i_n, i_d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
o_d = i_d * B + 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'))
if HAS_SCALE:
b_x = b_x * scale
b_m = tl.max(b_x, 0)
b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m
tl.store(z + i_n * tl.cdiv(D, B) + i_d, b_z)
def logsumexp_fwd(
x,
scale: Optional[float] = None,
dtype: Optional[torch.dtype] = None
):
r"""
Compute the logsumexp of the input tensor over the last dimension.
Args:
x (Tensor):
The input tensor of any shape.
scale (Optional[float]):
The scale applied to the input tensor. Default: `None`.
dtype (Optional[torch.dtype]):
The data type of the output tensor. Default: `None`.
Returns:
Tensor: The logsumexp of the input tensor.
"""
shape = x.shape
x = x.view(-1, shape[-1])
N, D = x.shape
B = min(triton.next_power_of_2(D), 64 * 1024)
ND = triton.cdiv(D, B)
z = x.new_empty(N, ND, dtype=torch.float)
logsumexp_fwd_kernel[(N, ND)](
x=x,
z=z,
scale=scale,
D=D,
B=B
)
z = z.logsumexp(-1).view(*shape[:-1])
if dtype is not None and dtype != torch.float:
z = z.to(dtype)
return z
|