Kernels:
Trusted publisher
File size: 3,053 Bytes
e19323e | 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 | # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
# Shared gate helpers reused across delta-rule family ops (KDA, GDN, ...).
import torch
import triton
import triton.language as tl
from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
@triton.jit
def fused_beta_sigmoid_fwd_kernel(
x,
y,
scale,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0).to(tl.int64)
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE).to(tl.int64)
mask = offs < n_elements
b_x = tl.load(x + offs, mask=mask, other=0).to(tl.float32)
b_y = scale * tl.sigmoid(b_x)
tl.store(y + offs, b_y.to(y.dtype.element_ty), mask=mask)
@triton.jit
def fused_beta_sigmoid_bwd_kernel(
x,
dy,
dx,
scale,
n_elements,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0).to(tl.int64)
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE).to(tl.int64)
mask = offs < n_elements
b_x = tl.load(x + offs, mask=mask, other=0).to(tl.float32)
b_dy = tl.load(dy + offs, mask=mask, other=0).to(tl.float32)
b_y = tl.sigmoid(b_x)
b_dx = b_dy * scale * b_y * (1.0 - b_y)
tl.store(dx + offs, b_dx.to(dx.dtype.element_ty), mask=mask)
_BETA_SIGMOID_BLOCK_SIZE = 2048
_BETA_SIGMOID_NUM_WARPS = 8
def fused_beta_sigmoid_fwd(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
y = torch.empty_like(x, dtype=torch.float32)
n_elements = x.numel()
grid = (triton.cdiv(n_elements, _BETA_SIGMOID_BLOCK_SIZE),)
fused_beta_sigmoid_fwd_kernel[grid](
x,
y,
scale,
n_elements,
BLOCK_SIZE=_BETA_SIGMOID_BLOCK_SIZE,
num_warps=_BETA_SIGMOID_NUM_WARPS,
)
return y
def fused_beta_sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
dx = torch.empty_like(x)
n_elements = x.numel()
grid = (triton.cdiv(n_elements, _BETA_SIGMOID_BLOCK_SIZE),)
fused_beta_sigmoid_bwd_kernel[grid](
x,
dy,
dx,
scale,
n_elements,
BLOCK_SIZE=_BETA_SIGMOID_BLOCK_SIZE,
num_warps=_BETA_SIGMOID_NUM_WARPS,
)
return dx
class BetaSigmoidFunction(torch.autograd.Function):
@staticmethod
@input_guard
@autocast_custom_fwd
def forward(ctx, x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
y = fused_beta_sigmoid_fwd(x, scale)
ctx.save_for_backward(x)
ctx.scale = scale
return y
@staticmethod
@input_guard
@autocast_custom_bwd
def backward(ctx, dy: torch.Tensor):
(x,) = ctx.saved_tensors
dx = fused_beta_sigmoid_bwd(x, dy, ctx.scale)
return dx.type_as(x), None
def fused_beta_sigmoid(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
return BetaSigmoidFunction.apply(x, scale)
|