| |
| |
|
|
| 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 |
|
|
