Spaces:
Running on Zero
Running on Zero
| from typing import Any | |
| import torch | |
| import triton | |
| import triton.language as tl | |
| from triton.language.extra.cuda import libdevice | |
| def _forward(X: tl.tensor, Y: tl.tensor, N: tl.tensor, OFFSET: tl.constexpr, CONSTANT: tl.constexpr, BLOCK: tl.constexpr) -> None: | |
| i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) | |
| x = tl.load(X + i, i < N, 0) | |
| # --------------------------------------------------------- | |
| # Evaluate softplus in FP32 with nonpositive exponential arguments. | |
| # Preserve input subtraction and output rounding in the input dtype. | |
| # --------------------------------------------------------- | |
| z = (x.to(tl.float32) - OFFSET).to(x.dtype).to(tl.float32) | |
| s = (tl.maximum(z, 0.0) + libdevice.log1p(tl.exp(-tl.abs(z)))).to(x.dtype) | |
| linear = (0.08 * x.to(tl.float32)).to(x.dtype) | |
| y = (s.to(tl.float32) - linear.to(tl.float32)).to(x.dtype) | |
| y = (y.to(tl.float32) - CONSTANT).to(x.dtype) | |
| tl.store(Y + i, y, i < N) | |
| def _backward(X: tl.tensor, G: tl.tensor, DX: tl.tensor, N: tl.tensor, OFFSET: tl.constexpr, BLOCK: tl.constexpr) -> None: | |
| i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) | |
| x = tl.load(X + i, i < N, 0) | |
| g = tl.load(G + i, i < N, 0) | |
| z = (x.to(tl.float32) - OFFSET).to(x.dtype).to(tl.float32) | |
| # --------------------------------------------------------- | |
| # Compute sigmoid with bounded exponentials for every input. | |
| # --------------------------------------------------------- | |
| sigmoid = tl.exp(tl.minimum(z, 0.0)) / (1.0 + tl.exp(-tl.abs(z))) | |
| nonlinear = (g.to(tl.float32) * sigmoid).to(x.dtype) | |
| linear = (-g.to(tl.float32) * 0.08).to(x.dtype) | |
| dx = nonlinear.to(tl.float32) + linear.to(tl.float32) | |
| tl.store(DX + i, dx, i < N) | |
| class _Swoosh(torch.autograd.Function): | |
| def forward(ctx: Any, x: torch.Tensor, offset: float, constant: float) -> torch.Tensor: | |
| x = x.contiguous() | |
| y = torch.empty_like(x) | |
| _forward[(triton.cdiv(x.numel(), 512),)](x, y, x.numel(), offset, constant, 512, enable_fp_fusion=False) | |
| ctx.save_for_backward(x) | |
| ctx.offset = offset | |
| return y | |
| def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None, None]: | |
| (x,) = ctx.saved_tensors | |
| dx = torch.empty_like(x) | |
| _backward[(triton.cdiv(x.numel(), 512),)](x, grad.contiguous(), dx, x.numel(), ctx.offset, 512, enable_fp_fusion=False) | |
| return dx, None, None | |
| def swoosh(x: torch.Tensor, activation: str) -> torch.Tensor: | |
| if not x.is_cuda: | |
| raise ValueError("The Triton Swoosh requires CUDA.") | |
| if activation not in ("SwooshL", "SwooshR"): | |
| raise ValueError(f"Unsupported activation: {activation}") | |
| return _Swoosh.apply(x, 4.0 if activation == "SwooshL" else 1.0, 0.035 if activation == "SwooshL" else 0.313261687) | |