mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
96 lines
3.8 KiB
96 lines
3.8 KiB
import torch
|
|
try:
|
|
import triton
|
|
import triton.language as tl
|
|
HAS_TRITON = True
|
|
except ImportError:
|
|
HAS_TRITON = False
|
|
print("please install triton from https://github.com/openai/triton")
|
|
|
|
if HAS_TRITON:
|
|
'''
|
|
softmax kernel is modified based on
|
|
https://github.com/openai/triton/blob/34817ecc954a6f4ca7b4dfb352fdde1f8bd49ca5/python/tutorials/02-fused-softmax.py
|
|
'''
|
|
@triton.jit
|
|
def softmax_kernel(output_ptr, input_ptr, row_stride, n_cols, mask_ptr, BLOCK_SIZE: tl.constexpr):
|
|
r""" the kernel function for implementing softmax operator
|
|
Args:
|
|
output_ptr: the output after finishing softmax operation, (N, hidden_dim)
|
|
input_ptr: the tensor of input, shape should be (N, hidden_dim)
|
|
n_cols(tl.constexpr): the number of cols of input
|
|
BLOCK_SIZE(tl.constexpr): the block_size of your hidden_dim dimension, typically BLOCK_SIZE >= hidden_dim
|
|
"""
|
|
row_idx = tl.program_id(0)
|
|
row_start_ptr = input_ptr + row_idx * row_stride
|
|
col_offsets = tl.arange(0, BLOCK_SIZE)
|
|
input_ptrs = row_start_ptr + col_offsets
|
|
row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf')).to(tl.float32)
|
|
row_minus_max = row - tl.max(row, axis=0)
|
|
|
|
if mask_ptr is not None:
|
|
# load mask into SRAM
|
|
mask_ptrs = (mask_ptr + (row_indx * row_stride)) + col_offsets
|
|
mask = tl.load(mask_ptrs, mask=col_offsets < n_cols, other=0).to(tl.float32)
|
|
|
|
# update
|
|
row_minus_max = row_minus_max + mask
|
|
|
|
numerator = tl.exp(row_minus_max)
|
|
denominator = tl.sum(numerator, axis=0)
|
|
softmax_output = numerator / denominator
|
|
output_row_start_ptr = output_ptr + row_idx * row_stride
|
|
output_ptrs = output_row_start_ptr + col_offsets
|
|
# Write back output to DRAM
|
|
tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols)
|
|
|
|
|
|
def softmax(input: torch.Tensor, mask: torch.Tensor = None, dim=-1) -> torch.Tensor:
|
|
if mask is not None:
|
|
assert input[-1] == mask[-1], "the last dimentions should be the same for input and mask"
|
|
assert dim == -1 or dim == len(input.shape)-1, "currently softmax layer only support last dimention"
|
|
|
|
hidden_dim = input.shape[-1]
|
|
output = torch.empty_like(input)
|
|
input = input.view(-1, hidden_dim)
|
|
if mask is not None:
|
|
mask = mask.view(-1, hidden_dim)
|
|
assert input.shape[0] == mask.shape[0], "the fist dimention of mask and input should be the same"
|
|
|
|
num_rows, num_cols = input.shape
|
|
block_size = max(triton.next_power_of_2(num_cols), 2)
|
|
num_warps = 16
|
|
if block_size >= 4096:
|
|
num_warps = 16
|
|
elif block_size >= 2048:
|
|
num_warps = 8
|
|
else:
|
|
num_warps = 4
|
|
|
|
if num_rows <= 350000:
|
|
grid = (num_rows,)
|
|
softmax_kernel[grid](output, input, input.stride(0), num_cols, mask, BLOCK_SIZE = block_size, num_warps=num_warps)
|
|
else:
|
|
grid = lambda meta: ()
|
|
|
|
grid = lambda meta: (
|
|
triton.cdiv(num_rows, meta["BLOCK_M"]),
|
|
)
|
|
|
|
BLOCK_M = 32
|
|
if block_size >= 4096:
|
|
BLOCK_M = 4
|
|
elif block_size >= 2048:
|
|
BLOCK_M = 8
|
|
|
|
softmax_kernel[grid](output_ptr = output,
|
|
input_ptr = input,
|
|
row_stride = input.stride(0),
|
|
n_rows = num_rows,
|
|
n_cols = num_cols,
|
|
mask_ptr = mask,
|
|
# currently manually setting up size
|
|
BLOCK_M = 32,
|
|
BLOCK_SIZE = block_size)
|
|
|
|
return output |