File size: 1,059 Bytes
2976eec | 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 | import torch
import triton
from ._triton_kernels.softmax import _softmax_kernel_online
from .utils.logger import AiterTritonLogger
_LOGGER = AiterTritonLogger()
def softmax(x):
"""
Computes row-wise softmax of a 2D input tensor.
Args:
x (torch.Tensor): Input tensor with shape (n_rows, n_cols). Must be on GPU.
Returns:
torch.Tensor: Output with same shape as x, softmax applied along last dimension.
"""
_LOGGER.info(f"SOFTMAX: x={tuple(x.shape)}")
n_rows, n_cols = x.shape
MAX_FUSED_SIZE = 65536 // x.element_size()
BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(n_cols))
y = torch.empty_like(x)
waves_per_eu = 2
num_warps = 8
num_stages = 2
num_programs = n_rows
grid = lambda meta: (num_programs,) # noqa: E731
_softmax_kernel_online[grid](
y,
x,
x.stride(0),
y.stride(0),
n_cols,
BLOCK_SIZE,
waves_per_eu=waves_per_eu,
num_warps=num_warps,
num_stages=num_stages,
)
return y
|