1093 lines
34 KiB
Python
1093 lines
34 KiB
Python
import logging
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import triton
|
|
|
|
from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8
|
|
from sglang.srt.utils import dispose_tensor, is_cuda
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_is_cuda = is_cuda()
|
|
if _is_cuda:
|
|
from sglang.srt.layers.quantization.fp8_kernel import (
|
|
sglang_per_token_group_quant_fp8 as per_token_group_quant_fp8,
|
|
)
|
|
|
|
try:
|
|
from deep_gemm import ceil_div
|
|
except ImportError:
|
|
logger.error(f"Failed to import ceil_div from deep_gemm.")
|
|
|
|
import triton.language as tl
|
|
|
|
|
|
@triton.jit
|
|
def deepep_permute_triton_kernel(
|
|
input_ptr,
|
|
gateup_input_ptr,
|
|
src2dst_ptr,
|
|
topk_ids_ptr,
|
|
a1_scales_ptr,
|
|
topk,
|
|
hidden_size,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
OutDtype = gateup_input_ptr.dtype.element_ty
|
|
|
|
src_idx = tl.program_id(0)
|
|
src2dst_ptr = src2dst_ptr + src_idx * topk
|
|
topk_ids_ptr = topk_ids_ptr + src_idx * topk
|
|
|
|
src_ptr = input_ptr + src_idx * hidden_size
|
|
|
|
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + tl.arange(0, BLOCK_SIZE)
|
|
mask = offset < hidden_size
|
|
in_data = tl.load(src_ptr + offset, mask=mask).to(OutDtype)
|
|
|
|
for idx in range(topk):
|
|
dst_idx = tl.load(src2dst_ptr + idx)
|
|
if dst_idx >= 0:
|
|
dst_ptr = gateup_input_ptr + dst_idx * hidden_size
|
|
tl.store(dst_ptr + offset, in_data, mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def deepep_post_reorder_triton_kernel(
|
|
down_output_ptr,
|
|
output_ptr,
|
|
src2dst_ptr,
|
|
topk_ids_ptr,
|
|
topk_weights_ptr,
|
|
topk,
|
|
hidden_size,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
InDtype = down_output_ptr.dtype.element_ty
|
|
|
|
src_idx = tl.program_id(0)
|
|
src2dst_ptr = src2dst_ptr + src_idx * topk
|
|
topk_ids_ptr = topk_ids_ptr + src_idx * topk
|
|
topk_weights_ptr = topk_weights_ptr + src_idx * topk
|
|
|
|
store_ptr = output_ptr + src_idx * hidden_size
|
|
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + tl.arange(0, BLOCK_SIZE)
|
|
mask = offset < hidden_size
|
|
sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype)
|
|
for idx in range(topk):
|
|
dst_idx = tl.load(src2dst_ptr + idx)
|
|
if dst_idx >= 0:
|
|
weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype)
|
|
load_ptr = down_output_ptr + dst_idx * hidden_size
|
|
in_data = tl.load(load_ptr + offset, mask=mask)
|
|
sum_vec += in_data * weigh_scale
|
|
tl.store(store_ptr + offset, sum_vec, mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def compute_src2dst_triton_kernel(
|
|
reorder_ids, src2dst, num_toks, BLOCK_SIZE: tl.constexpr
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
|
mask = dst_id < num_toks
|
|
src_id = tl.load(reorder_ids + dst_id, mask=mask)
|
|
tl.store(src2dst + src_id, dst_id, mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def deepep_compute_src2dst_triton_kernel(
|
|
reorder_ids, src2dst, num_toks, num_minus_one, BLOCK_SIZE: tl.constexpr
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
|
mask = dst_id < num_toks
|
|
src_id = tl.load(reorder_ids + dst_id, mask=mask)
|
|
num_invalid = tl.load(num_minus_one)
|
|
tl.store(src2dst + src_id, dst_id - num_invalid, mask=mask)
|
|
|
|
|
|
def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int):
|
|
reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)
|
|
seg_indptr = torch.empty(num_experts + 1, device=topk_ids.device, dtype=torch.int64)
|
|
src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int64)
|
|
|
|
# Find offset
|
|
expert_ids = torch.arange(
|
|
num_experts + 1, device=topk_ids.device, dtype=reorder_topk_ids.dtype
|
|
)
|
|
torch.searchsorted(reorder_topk_ids, expert_ids, out=seg_indptr)
|
|
num_minus_one = seg_indptr[0]
|
|
seg_indptr = seg_indptr - num_minus_one
|
|
|
|
BLOCK_SIZE = 512
|
|
grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),)
|
|
deepep_compute_src2dst_triton_kernel[grid](
|
|
reorder_ids, src2dst, topk_ids.numel(), num_minus_one, BLOCK_SIZE
|
|
)
|
|
reorder_topk_ids = reorder_topk_ids[num_minus_one:]
|
|
return reorder_topk_ids, src2dst, seg_indptr
|
|
|
|
|
|
@triton.jit
|
|
def compute_seg_indptr_triton_kernel(reorder_topk_ids, seg_indptr, num_toks):
|
|
expert = tl.program_id(0)
|
|
low = 0
|
|
high = num_toks - 1
|
|
target_location = -1
|
|
while low <= high:
|
|
mid = (low + high) // 2
|
|
|
|
if tl.load(reorder_topk_ids + mid) > expert:
|
|
high = mid - 1
|
|
else:
|
|
low = mid + 1
|
|
target_location = mid
|
|
tl.store(seg_indptr + expert + 1, target_location + 1)
|
|
|
|
|
|
def run_moe_ep_preproess(topk_ids: torch.Tensor, num_experts: int):
|
|
reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True)
|
|
seg_indptr = torch.zeros(num_experts + 1, device=topk_ids.device, dtype=torch.int64)
|
|
src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32)
|
|
|
|
compute_seg_indptr_triton_kernel[(num_experts,)](
|
|
reorder_topk_ids, seg_indptr, topk_ids.numel()
|
|
)
|
|
|
|
BLOCK_SIZE = 512
|
|
grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),)
|
|
compute_src2dst_triton_kernel[grid](
|
|
reorder_ids, src2dst, topk_ids.numel(), BLOCK_SIZE
|
|
)
|
|
return reorder_topk_ids, src2dst, seg_indptr
|
|
|
|
|
|
@triton.jit
|
|
def pre_reorder_triton_kernel(
|
|
input_ptr,
|
|
gateup_input_ptr,
|
|
src2dst_ptr,
|
|
topk_ids_ptr,
|
|
a1_scales_ptr,
|
|
start_expert_id,
|
|
end_expert_id,
|
|
topk,
|
|
hidden_size,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
use_per_token_if_dynamic: tl.constexpr,
|
|
):
|
|
OutDtype = gateup_input_ptr.dtype.element_ty
|
|
|
|
src_idx = tl.program_id(0)
|
|
src2dst_ptr = src2dst_ptr + src_idx * topk
|
|
topk_ids_ptr = topk_ids_ptr + src_idx * topk
|
|
src_ptr = input_ptr + src_idx * hidden_size
|
|
|
|
vec = tl.arange(0, BLOCK_SIZE)
|
|
|
|
if a1_scales_ptr is not None and use_per_token_if_dynamic:
|
|
scale = 1.0 / tl.load(a1_scales_ptr + src_idx)
|
|
|
|
for idx in range(topk):
|
|
expert_id = tl.load(topk_ids_ptr + idx)
|
|
if expert_id >= start_expert_id and expert_id <= end_expert_id:
|
|
if a1_scales_ptr is not None:
|
|
if not use_per_token_if_dynamic:
|
|
scale = 1.0 / tl.load(a1_scales_ptr + expert_id - start_expert_id)
|
|
else:
|
|
scale = 1.0
|
|
|
|
dst_idx = tl.load(src2dst_ptr + idx)
|
|
dst_ptr = gateup_input_ptr + dst_idx * hidden_size
|
|
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + vec
|
|
mask = offset < hidden_size
|
|
in_data = tl.load(src_ptr + offset, mask=mask).to(tl.float32)
|
|
out_data = (in_data * scale).to(OutDtype)
|
|
tl.store(dst_ptr + offset, out_data, mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def silu_and_mul_triton_kernel(
|
|
gateup_output,
|
|
down_input,
|
|
hidden_size,
|
|
reorder_topk_ids,
|
|
scales,
|
|
start_expert_id,
|
|
end_expert_id,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
InDtype = gateup_output.dtype.element_ty
|
|
OutDtype = down_input.dtype.element_ty
|
|
|
|
half_hidden_size = hidden_size // 2
|
|
|
|
pid = tl.program_id(0)
|
|
expert_id = tl.load(reorder_topk_ids + pid)
|
|
if expert_id >= start_expert_id and expert_id <= end_expert_id:
|
|
gateup_output_ptr = gateup_output + pid * hidden_size
|
|
gate_output_ptr = gateup_output_ptr
|
|
up_output_ptr = gateup_output_ptr + half_hidden_size
|
|
down_input_ptr = down_input + pid * half_hidden_size
|
|
|
|
if scales is not None:
|
|
scale = tl.load(scales + expert_id - start_expert_id)
|
|
scale = (1 / scale).to(InDtype)
|
|
else:
|
|
scale = 1
|
|
|
|
for start_offset in tl.range(0, half_hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + tl.arange(0, BLOCK_SIZE)
|
|
mask = offset < half_hidden_size
|
|
|
|
gate_output = tl.load(gate_output_ptr + offset, mask=mask).to(tl.float32)
|
|
up_output = tl.load(up_output_ptr + offset, mask=mask)
|
|
|
|
# silu & mul & quantize
|
|
gate_output = gate_output * tl.sigmoid(gate_output)
|
|
gate_output = gate_output.to(InDtype)
|
|
|
|
silu_mul_output = gate_output * up_output * scale
|
|
silu_mul_output = silu_mul_output.to(OutDtype)
|
|
tl.store(down_input_ptr + offset, silu_mul_output, mask=mask)
|
|
|
|
|
|
# copy from https://github.com/ModelTC/lightllm/blob/a000ab69098654df4731f5b12587dd4e7f0a4f41/lightllm/common/fused_moe/moe_silu_and_mul_mix_quant_ep.py
|
|
@triton.jit
|
|
def _silu_and_mul_post_quant_kernel(
|
|
input_ptr,
|
|
stride_input_0,
|
|
stride_input_1,
|
|
stride_input_2,
|
|
output_ptr,
|
|
stride_output_0,
|
|
stride_output_1,
|
|
stride_output_2,
|
|
output_scale_ptr,
|
|
stride_output_scale_0,
|
|
stride_output_scale_1,
|
|
stride_output_scale_2,
|
|
masked_m_ptr,
|
|
size_n,
|
|
fp8_max,
|
|
fp8_min,
|
|
BLOCK_N: tl.constexpr,
|
|
NUM_STAGE: tl.constexpr,
|
|
SCALE_UE8M0: tl.constexpr,
|
|
):
|
|
expert_id = tl.program_id(2)
|
|
token_id = tl.program_id(1)
|
|
hidden_dim_block_index = tl.program_id(0)
|
|
|
|
block_num_per_expert = tl.num_programs(1)
|
|
|
|
token_num_cur_expert = tl.load(masked_m_ptr + expert_id)
|
|
|
|
stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64)
|
|
stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64)
|
|
stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64)
|
|
stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64)
|
|
|
|
offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N)
|
|
input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d
|
|
output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d
|
|
output_scale_offs = (
|
|
output_scale_ptr
|
|
+ expert_id * stride_output_scale_0
|
|
+ hidden_dim_block_index * stride_output_scale_2
|
|
)
|
|
|
|
for token_index in tl.range(
|
|
token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE
|
|
):
|
|
gate = tl.load(
|
|
input_ptr_offs + token_index * stride_input_1,
|
|
mask=offs_in_d < size_n,
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
up = tl.load(
|
|
input_ptr_offs + token_index * stride_input_1 + size_n,
|
|
mask=offs_in_d < size_n,
|
|
other=0.0,
|
|
)
|
|
gate = gate / (1 + tl.exp(-gate))
|
|
gate = gate.to(input_ptr.dtype.element_ty)
|
|
gate_up = up * gate
|
|
_absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10)
|
|
output_s = _absmax / fp8_max
|
|
if SCALE_UE8M0:
|
|
output_s = tl.exp2(tl.ceil(tl.log2(tl.abs(output_s))))
|
|
output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to(
|
|
output_ptr.dtype.element_ty
|
|
)
|
|
tl.store(
|
|
output_ptr_offs + token_index * stride_output_1,
|
|
output_q,
|
|
mask=offs_in_d < size_n,
|
|
)
|
|
tl.store(
|
|
output_scale_offs + token_index * stride_output_scale_1,
|
|
output_s,
|
|
)
|
|
|
|
|
|
def silu_and_mul_masked_post_quant_fwd(
|
|
input: torch.Tensor,
|
|
output: torch.Tensor,
|
|
output_scale: torch.Tensor,
|
|
quant_group_size: int,
|
|
masked_m: torch.Tensor,
|
|
scale_ue8m0: bool = False,
|
|
):
|
|
"""
|
|
input shape [expert_num, token_num_padded, hidden_dim]
|
|
output shape [expert_num, token_num_padded, hidden_dim // 2], dtype fp8
|
|
output_scale [expert_num token_num_paddded, hidden_dim // 2 // 128] dtype float32
|
|
quant_group_size int,
|
|
masked_m shape [expert_num],
|
|
"""
|
|
|
|
assert input.is_contiguous()
|
|
assert output.dtype == torch.float8_e4m3fn
|
|
assert output.is_contiguous()
|
|
assert len(input.shape) == 3
|
|
assert input.shape[0] == masked_m.shape[0]
|
|
assert input.shape[-1] % 2 == 0
|
|
|
|
size_n = input.shape[-1] // 2
|
|
assert size_n % quant_group_size == 0
|
|
|
|
expert_num = len(masked_m)
|
|
|
|
if expert_num < 4:
|
|
BLOCK_NUM_PER_EXPERT = 64
|
|
else:
|
|
BLOCK_NUM_PER_EXPERT = 32
|
|
|
|
BLOCK_N = quant_group_size
|
|
num_warps = 1
|
|
NUM_STAGES = 6
|
|
hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N)
|
|
assert BLOCK_N % quant_group_size == 0
|
|
|
|
grid = (
|
|
hidden_dim_split_block_num,
|
|
BLOCK_NUM_PER_EXPERT,
|
|
expert_num,
|
|
)
|
|
|
|
finfo = torch.finfo(torch.float8_e4m3fn)
|
|
fp8_max = finfo.max
|
|
fp8_min = -fp8_max
|
|
|
|
_silu_and_mul_post_quant_kernel[grid](
|
|
input,
|
|
*input.stride(),
|
|
output,
|
|
*output.stride(),
|
|
output_scale,
|
|
*output_scale.stride(),
|
|
masked_m,
|
|
size_n,
|
|
fp8_max,
|
|
fp8_min,
|
|
BLOCK_N=BLOCK_N,
|
|
NUM_STAGE=NUM_STAGES,
|
|
num_warps=num_warps,
|
|
SCALE_UE8M0=scale_ue8m0,
|
|
)
|
|
return
|
|
|
|
|
|
@triton.jit
|
|
def tanh(x):
|
|
return 2 * tl.sigmoid(2 * x) - 1
|
|
|
|
|
|
@triton.jit
|
|
def gelu_and_mul_triton_kernel(
|
|
gateup_output,
|
|
down_input,
|
|
hidden_size,
|
|
reorder_topk_ids,
|
|
scales,
|
|
start_expert_id,
|
|
end_expert_id,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
InDtype = gateup_output.dtype.element_ty
|
|
OutDtype = down_input.dtype.element_ty
|
|
|
|
half_hidden_size = hidden_size // 2
|
|
|
|
pid = tl.program_id(0)
|
|
expert_id = tl.load(reorder_topk_ids + pid)
|
|
if expert_id >= start_expert_id and expert_id <= end_expert_id:
|
|
gateup_output_ptr = gateup_output + pid * hidden_size
|
|
gate_output_ptr = gateup_output_ptr
|
|
up_output_ptr = gateup_output_ptr + half_hidden_size
|
|
down_input_ptr = down_input + pid * half_hidden_size
|
|
|
|
if scales is not None:
|
|
scale = tl.load(scales + expert_id - start_expert_id)
|
|
scale = (1 / scale).to(InDtype)
|
|
else:
|
|
scale = 1
|
|
|
|
for start_offset in tl.range(0, half_hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + tl.arange(0, BLOCK_SIZE)
|
|
mask = offset < half_hidden_size
|
|
|
|
gate_output = tl.load(gate_output_ptr + offset, mask=mask).to(tl.float32)
|
|
up_output = tl.load(up_output_ptr + offset, mask=mask)
|
|
|
|
# gelu & mul & quantize
|
|
# https://pytorch.org/docs/stable/generated/torch.nn.GELU.html
|
|
# sqrt(2/pi)
|
|
kAlpha = 0.7978845608028654
|
|
gate_output = (
|
|
0.5
|
|
* gate_output
|
|
* (
|
|
1
|
|
+ tanh(
|
|
kAlpha
|
|
* (
|
|
gate_output
|
|
+ 0.044715 * gate_output * gate_output * gate_output
|
|
)
|
|
)
|
|
)
|
|
)
|
|
gate_output = gate_output.to(InDtype)
|
|
|
|
gelu_mul_output = gate_output * up_output * scale
|
|
gelu_mul_output = gelu_mul_output.to(OutDtype)
|
|
tl.store(down_input_ptr + offset, gelu_mul_output, mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def post_reorder_triton_kernel(
|
|
down_output_ptr,
|
|
output_ptr,
|
|
src2dst_ptr,
|
|
topk_ids_ptr,
|
|
topk_weights_ptr,
|
|
start_expert_id,
|
|
end_expert_id,
|
|
topk,
|
|
hidden_size,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
InDtype = down_output_ptr.dtype.element_ty
|
|
|
|
src_idx = tl.program_id(0)
|
|
src2dst_ptr = src2dst_ptr + src_idx * topk
|
|
topk_ids_ptr = topk_ids_ptr + src_idx * topk
|
|
topk_weights_ptr = topk_weights_ptr + src_idx * topk
|
|
|
|
computed = False
|
|
store_ptr = output_ptr + src_idx * hidden_size
|
|
|
|
vec = tl.arange(0, BLOCK_SIZE)
|
|
|
|
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + vec
|
|
mask = offset < hidden_size
|
|
|
|
sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype)
|
|
for idx in range(topk):
|
|
expert_id = tl.load(topk_ids_ptr + idx)
|
|
if expert_id >= start_expert_id and expert_id <= end_expert_id:
|
|
computed = True
|
|
dst_idx = tl.load(src2dst_ptr + idx)
|
|
weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype)
|
|
load_ptr = down_output_ptr + dst_idx * hidden_size
|
|
in_data = tl.load(load_ptr + offset, mask=mask)
|
|
sum_vec += in_data * weigh_scale
|
|
tl.store(store_ptr + offset, sum_vec, mask=mask)
|
|
|
|
if computed == False:
|
|
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
|
|
offset = start_offset + vec
|
|
mask = offset < hidden_size
|
|
tl.store(
|
|
store_ptr + offset, tl.zeros([BLOCK_SIZE], dtype=InDtype), mask=mask
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def compute_m_range(
|
|
pid,
|
|
batch_size,
|
|
seg_indptr,
|
|
weight_indices,
|
|
m_num_tiles_indptr,
|
|
BLOCK_SIZE_M: tl.constexpr,
|
|
):
|
|
idx = 0
|
|
for bs in range(batch_size):
|
|
tiles = tl.load(m_num_tiles_indptr + bs)
|
|
if pid >= tiles:
|
|
idx = bs
|
|
|
|
idx_start = tl.load(m_num_tiles_indptr + idx)
|
|
|
|
m_range_start = tl.load(seg_indptr + idx) + (pid - idx_start) * BLOCK_SIZE_M
|
|
m_range_end = min(tl.load(seg_indptr + idx + 1), m_range_start + BLOCK_SIZE_M)
|
|
expert_id = tl.load(weight_indices + idx)
|
|
return m_range_start, m_range_end, expert_id
|
|
|
|
|
|
@triton.jit
|
|
def grouped_gemm_triton_kernel(
|
|
a,
|
|
b,
|
|
c,
|
|
batch_size,
|
|
N,
|
|
K,
|
|
seg_indptr,
|
|
weight_indices,
|
|
m_num_tiles_indptr,
|
|
scale_a,
|
|
scale_b,
|
|
use_fp8_w8a8: tl.constexpr,
|
|
group_n: tl.constexpr,
|
|
group_k: tl.constexpr,
|
|
a_stride_0: tl.constexpr,
|
|
b_stride_0: tl.constexpr,
|
|
b_stride_1: tl.constexpr,
|
|
as_stride_0: tl.constexpr,
|
|
as_stride_1: tl.constexpr,
|
|
bs_stride_0: tl.constexpr,
|
|
bs_stride_2: tl.constexpr,
|
|
bs_stride_1: tl.constexpr,
|
|
use_per_token_if_dynamic: tl.constexpr,
|
|
BLOCK_SIZE_M: tl.constexpr,
|
|
BLOCK_SIZE_N: tl.constexpr,
|
|
BLOCK_SIZE_K: tl.constexpr,
|
|
):
|
|
c_dtype = c.dtype.element_ty
|
|
|
|
pid_m = tl.program_id(0)
|
|
pid_n = tl.program_id(1)
|
|
total_m_block = tl.load(m_num_tiles_indptr + batch_size)
|
|
if pid_m >= total_m_block:
|
|
return
|
|
|
|
m_range_start, m_range_end, expert_id = compute_m_range(
|
|
pid_m, batch_size, seg_indptr, weight_indices, m_num_tiles_indptr, BLOCK_SIZE_M
|
|
)
|
|
if m_range_end - m_range_start == 0:
|
|
return
|
|
|
|
n_range_start = pid_n * BLOCK_SIZE_N
|
|
n_range_end = min(n_range_start + BLOCK_SIZE_N, N)
|
|
|
|
offs_am = tl.arange(0, BLOCK_SIZE_M)
|
|
offs_bn = tl.arange(0, BLOCK_SIZE_N)
|
|
|
|
offs_am = tl.where(offs_am < m_range_end - m_range_start, offs_am, 0)
|
|
offs_bn = tl.where(offs_bn < n_range_end - n_range_start, offs_bn, 0)
|
|
offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M)
|
|
offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N)
|
|
offs_k = tl.arange(0, BLOCK_SIZE_K)
|
|
|
|
a_ptr = a + (m_range_start + offs_am[:, None]) * a_stride_0 + offs_k[None, :]
|
|
b_ptr = b + (
|
|
(expert_id * b_stride_0)
|
|
+ (n_range_start + offs_bn[:, None]) * b_stride_1
|
|
+ offs_k[None, :]
|
|
)
|
|
|
|
if group_k > 0 and group_n > 0:
|
|
a_scale_ptrs = scale_a + (m_range_start + offs_am[:, None]) * as_stride_0
|
|
offs_bsn = (n_range_start + offs_bn) // group_n
|
|
b_scale_ptrs = scale_b + (expert_id * bs_stride_0) + offs_bsn * bs_stride_1
|
|
|
|
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
|
|
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
|
|
a_tile = tl.load(
|
|
a_ptr, mask=offs_k[None, :] < (K - k * BLOCK_SIZE_K), other=0.0
|
|
)
|
|
b_tile = tl.load(
|
|
b_ptr, mask=offs_k[None, :] < (K - k * BLOCK_SIZE_K), other=0.0
|
|
)
|
|
|
|
if group_k > 0 and group_n > 0:
|
|
k_start = k * BLOCK_SIZE_K
|
|
offs_ks = k_start // group_k
|
|
a_scale = tl.load(a_scale_ptrs + offs_ks * as_stride_1)
|
|
b_scale = tl.load(b_scale_ptrs + offs_ks * bs_stride_2)
|
|
accumulator += tl.dot(a_tile, b_tile.T) * a_scale * b_scale[None, :]
|
|
else:
|
|
accumulator = tl.dot(a_tile, b_tile.T, accumulator)
|
|
a_ptr += BLOCK_SIZE_K
|
|
b_ptr += BLOCK_SIZE_K
|
|
|
|
if use_fp8_w8a8 and not (group_k > 0 and group_n > 0):
|
|
if use_per_token_if_dynamic:
|
|
scale_a_value = tl.load(scale_a + (m_range_start + offs_am[:, None]))
|
|
else:
|
|
scale_a_value = tl.load(scale_a + expert_id)
|
|
scale_b_value = tl.load(scale_b + expert_id)
|
|
accumulator *= scale_a_value * scale_b_value
|
|
|
|
c_tile = accumulator.to(c_dtype)
|
|
|
|
offs_cm = m_range_start + tl.arange(0, BLOCK_SIZE_M)
|
|
offs_cn = n_range_start + tl.arange(0, BLOCK_SIZE_N)
|
|
c_ptr = c + offs_cm[:, None] * N + offs_cn[None, :]
|
|
c_mask = (offs_cm[:, None] < m_range_end) & (offs_cn[None, :] < n_range_end)
|
|
tl.store(c_ptr, c_tile, mask=c_mask)
|
|
|
|
|
|
@triton.jit
|
|
def compute_m_num_tiles_indptr(
|
|
m_num_tiles_indptr, seg_indptr, batch_size: tl.constexpr, BLOCK_SIZE_M: tl.constexpr
|
|
):
|
|
for bs in range(batch_size):
|
|
m = tl.load(seg_indptr + bs + 1) - tl.load(seg_indptr + bs)
|
|
cur_num_tiles = tl.cdiv(m, BLOCK_SIZE_M)
|
|
pre_num_tiles = tl.load(m_num_tiles_indptr + bs)
|
|
tl.store(m_num_tiles_indptr + bs + 1, pre_num_tiles + cur_num_tiles)
|
|
|
|
|
|
def grouped_gemm_triton(
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
c: torch.Tensor,
|
|
batch_size: int,
|
|
weight_column_major: bool,
|
|
seg_indptr: Optional[torch.Tensor] = None,
|
|
weight_indices: Optional[torch.Tensor] = None,
|
|
use_fp8_w8a8: bool = False,
|
|
scale_a: torch.Tensor = None,
|
|
scale_b: torch.Tensor = None,
|
|
block_shape: Optional[List[int]] = None,
|
|
c_dtype=None,
|
|
use_per_token_if_dynamic: bool = True,
|
|
):
|
|
assert weight_column_major == True # TODO: more
|
|
if use_fp8_w8a8 and block_shape is None:
|
|
assert scale_a is not None and scale_b is not None
|
|
|
|
if block_shape is not None:
|
|
a_original = a
|
|
|
|
assert len(block_shape) == 2
|
|
block_n, block_k = block_shape[0], block_shape[1]
|
|
a, scale_a = per_token_group_quant_fp8(a, block_k)
|
|
|
|
assert triton.cdiv(a.shape[-1], block_k) == scale_a.shape[-1]
|
|
assert triton.cdiv(b.shape[-2], block_n) == scale_b.shape[-2]
|
|
assert triton.cdiv(b.shape[-1], block_k) == scale_b.shape[-1]
|
|
|
|
dispose_tensor(a_original)
|
|
|
|
# TODO: adjust config or tune kernel
|
|
# Reduce block size to prevent L40 shared memory overflow.
|
|
config = {
|
|
"BLOCK_SIZE_M": 64,
|
|
"BLOCK_SIZE_N": 32,
|
|
"BLOCK_SIZE_K": 128,
|
|
}
|
|
|
|
m_num_tiles_indptr = torch.zeros(batch_size + 1, device=a.device, dtype=torch.int64)
|
|
compute_m_num_tiles_indptr[(1,)](
|
|
m_num_tiles_indptr, seg_indptr, batch_size, config["BLOCK_SIZE_M"]
|
|
)
|
|
|
|
if c is None:
|
|
assert c_dtype is not None
|
|
c = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=c_dtype)
|
|
|
|
grid = lambda META: (
|
|
triton.cdiv(a.size(0), META["BLOCK_SIZE_M"]) + batch_size,
|
|
triton.cdiv(b.size(1), META["BLOCK_SIZE_N"]),
|
|
)
|
|
|
|
if use_fp8_w8a8 and block_shape is None and use_per_token_if_dynamic:
|
|
assert (
|
|
scale_a.shape[0] == a.shape[0]
|
|
), f"scale_a.shape: {scale_a.shape}, a.shape: {a.shape}"
|
|
|
|
grouped_gemm_triton_kernel[grid](
|
|
a,
|
|
b,
|
|
c,
|
|
batch_size,
|
|
b.size(1),
|
|
b.size(2),
|
|
seg_indptr,
|
|
weight_indices,
|
|
m_num_tiles_indptr,
|
|
scale_a,
|
|
scale_b,
|
|
use_fp8_w8a8,
|
|
0 if block_shape is None else block_shape[0],
|
|
0 if block_shape is None else block_shape[1],
|
|
a.stride(0),
|
|
b.stride(0),
|
|
b.stride(1),
|
|
scale_a.stride(0) if scale_a is not None and scale_a.ndim == 2 else 0,
|
|
scale_a.stride(1) if scale_a is not None and scale_a.ndim == 2 else 0,
|
|
scale_b.stride(0) if scale_b is not None and scale_b.ndim >= 2 else 0,
|
|
scale_b.stride(2) if scale_b is not None and scale_b.ndim == 3 else 0,
|
|
scale_b.stride(1) if scale_b is not None and scale_b.ndim >= 2 else 0,
|
|
use_per_token_if_dynamic,
|
|
**config,
|
|
)
|
|
return c
|
|
|
|
|
|
@triton.jit
|
|
def _fwd_kernel_ep_scatter_1(
|
|
num_recv_tokens_per_expert,
|
|
expert_start_loc,
|
|
m_indices,
|
|
num_experts: tl.constexpr,
|
|
BLOCK_E: tl.constexpr,
|
|
BLOCK_EXPERT_NUM: tl.constexpr,
|
|
):
|
|
cur_expert = tl.program_id(0)
|
|
|
|
offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM)
|
|
tokens_per_expert = tl.load(
|
|
num_recv_tokens_per_expert + offset_cumsum,
|
|
mask=offset_cumsum < num_experts,
|
|
other=0,
|
|
)
|
|
cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
|
|
tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts)
|
|
|
|
cur_expert_start = tl.load(expert_start_loc + cur_expert)
|
|
cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)
|
|
|
|
m_indices_start_ptr = m_indices + cur_expert_start
|
|
off_expert = tl.arange(0, BLOCK_E)
|
|
|
|
for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4):
|
|
tl.store(
|
|
m_indices_start_ptr + start_m + off_expert,
|
|
cur_expert,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _fwd_kernel_ep_scatter_2(
|
|
total_token_num,
|
|
expert_start_loc,
|
|
recv_x,
|
|
recv_x_stride0,
|
|
recv_x_stride1,
|
|
recv_x_scale,
|
|
recv_x_scale_stride0,
|
|
recv_x_scale_stride1,
|
|
recv_topk,
|
|
recv_topk_stride0,
|
|
recv_topk_stride1,
|
|
output_tensor,
|
|
output_tensor_stride0,
|
|
output_tensor_stride1,
|
|
output_tensor_scale,
|
|
output_tensor_scale_stride0,
|
|
output_tensor_scale_stride1,
|
|
output_index,
|
|
output_index_stride0,
|
|
output_index_stride1,
|
|
topk_num: tl.constexpr,
|
|
HIDDEN_SIZE: tl.constexpr,
|
|
HIDDEN_SIZE_PAD: tl.constexpr,
|
|
SCALE_HIDDEN_SIZE: tl.constexpr,
|
|
SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
|
|
):
|
|
start_token_id = tl.program_id(0)
|
|
grid_num = tl.num_programs(0)
|
|
|
|
offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
|
|
mask = offset_in < HIDDEN_SIZE
|
|
|
|
offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
|
|
mask_s = offset_in_s < SCALE_HIDDEN_SIZE
|
|
|
|
for token_id_int32 in range(start_token_id, total_token_num, grid_num):
|
|
token_id = token_id_int32.to(tl.int64)
|
|
to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
|
|
to_copy_s = tl.load(
|
|
recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s
|
|
)
|
|
|
|
for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4):
|
|
topk_index = topk_idx_int32.to(tl.int64)
|
|
expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
|
|
if expert_id >= 0:
|
|
dest_token_index_int32 = tl.atomic_add(expert_start_loc + expert_id, 1)
|
|
dest_token_index = dest_token_index_int32.to(tl.int64)
|
|
|
|
tl.store(
|
|
output_index + token_id * output_index_stride0 + topk_index,
|
|
dest_token_index_int32,
|
|
)
|
|
output_tensor_ptr = (
|
|
output_tensor + dest_token_index * output_tensor_stride0
|
|
)
|
|
output_tensor_scale_ptr = (
|
|
output_tensor_scale + dest_token_index * output_tensor_scale_stride0
|
|
)
|
|
tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
|
|
tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s)
|
|
|
|
|
|
# copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/deepep_scatter_gather.py
|
|
@torch.no_grad()
|
|
def ep_scatter(
|
|
recv_x: torch.Tensor,
|
|
recv_x_scale: torch.Tensor,
|
|
recv_topk: torch.Tensor,
|
|
num_recv_tokens_per_expert: torch.Tensor,
|
|
expert_start_loc: torch.Tensor,
|
|
output_tensor: torch.Tensor,
|
|
output_tensor_scale: torch.Tensor,
|
|
m_indices: torch.Tensor,
|
|
output_index: torch.Tensor,
|
|
):
|
|
BLOCK_E = 128 # token num of per expert is aligned to 128
|
|
BLOCK_D = 128 # block size of quantization
|
|
num_warps = 8
|
|
num_experts = num_recv_tokens_per_expert.shape[0]
|
|
hidden_size = recv_x.shape[1]
|
|
# grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
|
|
grid = num_experts
|
|
|
|
assert m_indices.shape[0] % BLOCK_E == 0
|
|
|
|
_fwd_kernel_ep_scatter_1[(grid,)](
|
|
num_recv_tokens_per_expert,
|
|
expert_start_loc,
|
|
m_indices,
|
|
num_experts=num_experts,
|
|
num_warps=num_warps,
|
|
BLOCK_E=BLOCK_E,
|
|
BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
|
|
)
|
|
|
|
grid = min(recv_topk.shape[0], 1024 * 8)
|
|
|
|
_fwd_kernel_ep_scatter_2[(grid,)](
|
|
recv_topk.shape[0],
|
|
expert_start_loc,
|
|
recv_x,
|
|
recv_x.stride(0),
|
|
recv_x.stride(1),
|
|
recv_x_scale,
|
|
recv_x_scale.stride(0),
|
|
recv_x_scale.stride(1),
|
|
recv_topk,
|
|
recv_topk.stride(0),
|
|
recv_topk.stride(1),
|
|
output_tensor,
|
|
output_tensor.stride(0),
|
|
output_tensor.stride(1),
|
|
output_tensor_scale,
|
|
output_tensor_scale.stride(0),
|
|
output_tensor_scale.stride(1),
|
|
output_index,
|
|
output_index.stride(0),
|
|
output_index.stride(1),
|
|
topk_num=recv_topk.shape[1],
|
|
num_warps=num_warps,
|
|
HIDDEN_SIZE=hidden_size,
|
|
HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
|
|
SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D,
|
|
SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D),
|
|
)
|
|
return
|
|
|
|
|
|
@triton.jit
|
|
def _fwd_kernel_ep_gather(
|
|
total_token_num,
|
|
input_tensor,
|
|
input_tensor_stride0,
|
|
input_tensor_stride1,
|
|
recv_topk_ids,
|
|
recv_topk_ids_stride0,
|
|
recv_topk_ids_stride1,
|
|
recv_topk_weight,
|
|
recv_topk_weight_stride0,
|
|
recv_topk_weight_stride1,
|
|
input_index,
|
|
input_index_stride0,
|
|
input_index_stride1,
|
|
output_tensor,
|
|
output_tensor_stride0,
|
|
output_tensor_stride1,
|
|
topk_num: tl.constexpr,
|
|
BLOCK_D: tl.constexpr,
|
|
):
|
|
cur_block_int32 = tl.program_id(0)
|
|
cur_block = cur_block_int32.to(tl.int64)
|
|
|
|
start_cur_token_int32 = tl.program_id(1)
|
|
|
|
grid_num = tl.num_programs(1)
|
|
|
|
for cur_token_int32 in range(start_cur_token_int32, total_token_num, grid_num):
|
|
cur_token = cur_token_int32.to(tl.int64)
|
|
|
|
off_d = tl.arange(0, BLOCK_D)
|
|
accumulator = tl.zeros([BLOCK_D], dtype=tl.float32)
|
|
|
|
for topk_index_int32 in range(0, topk_num):
|
|
topk_index = topk_index_int32.to(tl.int64)
|
|
|
|
expert_id = tl.load(
|
|
recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index
|
|
)
|
|
if expert_id >= 0:
|
|
source_token_index_int32 = tl.load(
|
|
input_index + cur_token * input_index_stride0 + topk_index
|
|
)
|
|
source_token_index = source_token_index_int32.to(tl.int64)
|
|
|
|
acc_weight = tl.load(
|
|
recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index
|
|
)
|
|
tmp = tl.load(
|
|
input_tensor
|
|
+ source_token_index * input_tensor_stride0
|
|
+ cur_block * BLOCK_D
|
|
+ off_d
|
|
)
|
|
accumulator += tmp.to(tl.float32) * acc_weight
|
|
|
|
tl.store(
|
|
output_tensor
|
|
+ cur_token * output_tensor_stride0
|
|
+ cur_block * BLOCK_D
|
|
+ off_d,
|
|
accumulator.to(output_tensor.dtype.element_ty),
|
|
)
|
|
|
|
|
|
@torch.no_grad()
|
|
def ep_gather(
|
|
input_tensor: torch.Tensor,
|
|
recv_topk_ids: torch.Tensor,
|
|
recv_topk_weight: torch.Tensor,
|
|
input_index: torch.Tensor,
|
|
output_tensor: torch.Tensor,
|
|
):
|
|
BLOCK_D = 1024 # block size of quantization
|
|
num_warps = 2
|
|
num_tokens = output_tensor.shape[0]
|
|
hidden_size = input_tensor.shape[1]
|
|
assert hidden_size % BLOCK_D == 0
|
|
grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
|
|
_fwd_kernel_ep_gather[grid](
|
|
num_tokens,
|
|
input_tensor,
|
|
input_tensor.stride(0),
|
|
input_tensor.stride(1),
|
|
recv_topk_ids,
|
|
recv_topk_ids.stride(0),
|
|
recv_topk_ids.stride(1),
|
|
recv_topk_weight,
|
|
recv_topk_weight.stride(0),
|
|
recv_topk_weight.stride(1),
|
|
input_index,
|
|
input_index.stride(0),
|
|
input_index.stride(1),
|
|
output_tensor,
|
|
output_tensor.stride(0),
|
|
output_tensor.stride(1),
|
|
topk_num=recv_topk_ids.shape[1],
|
|
num_warps=num_warps,
|
|
BLOCK_D=BLOCK_D,
|
|
)
|
|
return
|
|
|
|
|
|
# copy from
|
|
# https://github.com/deepseek-ai/DeepGEMM/blob/bd2a77552886b98c205af12f8d7d2d61247c4b27/deep_gemm/jit_kernels/utils.py#L58
|
|
def get_tma_aligned_size(x: int, element_size: int) -> int:
|
|
"""
|
|
Global memory address of TMA must be 16-byte aligned.
|
|
Since we use column-major layout for the LHS scaling tensor,
|
|
the M-axis of the LHS scaling tensor needs to be padded to a multiple of 16 bytes.
|
|
|
|
Arguments:
|
|
x: original M-axis shape of the LHS scaling tensor.
|
|
element_size: element size of the LHS scaling tensor.
|
|
|
|
Returns:
|
|
M-axis shape of the LHS scaling tensor after padding.
|
|
"""
|
|
tma_alignment_bytes = 16
|
|
assert tma_alignment_bytes % element_size == 0
|
|
alignment = tma_alignment_bytes // element_size
|
|
return ceil_div(x, alignment) * alignment
|
|
|
|
|
|
@triton.jit
|
|
def _tma_align_input_scale_kernel(
|
|
input_scale_ptr,
|
|
output_ptr,
|
|
m,
|
|
k_div_block_size,
|
|
input_scale_stride_m,
|
|
input_scale_stride_k,
|
|
output_stride_m,
|
|
output_stride_k,
|
|
BLOCK_SIZE_K: tl.constexpr,
|
|
):
|
|
pid_m = tl.program_id(axis=0)
|
|
grid_m = tl.num_programs(0)
|
|
k_offsets = tl.arange(0, BLOCK_SIZE_K)
|
|
|
|
for m_base in range(pid_m, m, grid_m):
|
|
input_offset = (
|
|
input_scale_ptr
|
|
+ m_base * input_scale_stride_m
|
|
+ k_offsets * input_scale_stride_k
|
|
)
|
|
input_data = tl.load(input_offset, mask=k_offsets < k_div_block_size)
|
|
|
|
output_offset = (
|
|
output_ptr + k_offsets * output_stride_k + m_base * output_stride_m
|
|
)
|
|
tl.store(output_offset, input_data, mask=k_offsets < k_div_block_size)
|
|
|
|
|
|
# copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/quantization/triton_quant/fp8/fp8act_quant_kernel.py
|
|
def tma_align_input_scale(input_scale: torch.Tensor):
|
|
assert input_scale.dim() == 2
|
|
m, k_div_block_size = input_scale.shape
|
|
padd_m = get_tma_aligned_size(m, input_scale.element_size())
|
|
output = torch.empty(
|
|
(k_div_block_size, padd_m), dtype=input_scale.dtype, device=input_scale.device
|
|
)
|
|
|
|
grid_m = min(m, 8192)
|
|
BLOCK_SIZE_K = triton.next_power_of_2(k_div_block_size)
|
|
|
|
_tma_align_input_scale_kernel[(grid_m,)](
|
|
input_scale_ptr=input_scale,
|
|
output_ptr=output,
|
|
m=m,
|
|
k_div_block_size=k_div_block_size,
|
|
input_scale_stride_m=input_scale.stride(0),
|
|
input_scale_stride_k=input_scale.stride(1),
|
|
output_stride_m=output.stride(1), # Note: these are swapped
|
|
output_stride_k=output.stride(0), # for column-major
|
|
BLOCK_SIZE_K=BLOCK_SIZE_K,
|
|
)
|
|
return output.t()[:m]
|