Spaces:
Sleeping
Sleeping
File size: 1,459 Bytes
82e11f0 | 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 | from typing import Any
import torch
import triton
import triton.language as tl
@triton.jit(do_not_specialize=["N"])
def _backward(Y: tl.tensor, G: tl.tensor, DX: tl.tensor, N: tl.tensor, BLOCK: tl.constexpr) -> None:
row = tl.program_id(0)
col = tl.arange(0, BLOCK)
y = tl.load(Y + row * N + col, col < N, 0).to(tl.float32)
g = tl.load(G + row * N + col, col < N, 0).to(tl.float32)
product = y * g
dx = product - y * tl.sum(product, 0)
tl.store(DX + row * N + col, dx, col < N)
class _Softmax(torch.autograd.Function):
@staticmethod
def forward(ctx: Any, x: torch.Tensor) -> torch.Tensor:
y = x.softmax(dim=-1)
if x.requires_grad and torch.is_autocast_enabled("cuda"):
y = y.to(torch.get_autocast_dtype("cuda"))
ctx.save_for_backward(y)
ctx.input_dtype = x.dtype
return y
@staticmethod
def backward(ctx: Any, grad: torch.Tensor) -> torch.Tensor:
(y,) = ctx.saved_tensors
grad = grad.contiguous()
dx = torch.empty(y.shape, device=y.device, dtype=ctx.input_dtype)
n = y.shape[-1]
_backward[(y.numel() // n,)](y, grad, dx, n, triton.next_power_of_2(n), enable_fp_fusion=False)
return dx
def softmax(x: torch.Tensor, dim: int) -> torch.Tensor:
if not x.is_cuda or dim not in (-1, x.ndim - 1):
raise ValueError("The Triton softmax requires CUDA and the last dimension.")
return _Softmax.apply(x)
|