wehub-resource-sync
769 lines
33 KiB
Python
769 lines
33 KiB
Python
"""Mixture of Experts operators"""
|
|
|
|
from typing import Literal, Optional, Tuple # noqa: UP035
|
|
|
|
from tvm import DataType, DataTypeCode, s_tir, tirx
|
|
from tvm.relax.frontend.nn import Tensor, op
|
|
from tvm.script import tirx as T
|
|
|
|
# mypy: disable-error-code="attr-defined,valid-type,name-defined"
|
|
|
|
|
|
def gemv(x: Tensor, w: Tensor, indptr: Tensor) -> Tensor:
|
|
"""GEMV for project-in (e1-e3) or project-out (e2) in MLP.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
For project-in, the input tensor of shape (1, in_features); and for project-out, the input
|
|
shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
|
|
experts per token.
|
|
|
|
w : Tensor
|
|
The weight tensor of shape (local_experts, out_features, in_features), where `local_experts`
|
|
is the total number of experts.
|
|
|
|
indptr : Tensor
|
|
The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
|
|
number of activated experts per token.
|
|
|
|
Returns
|
|
-------
|
|
out : Tensor
|
|
The output tensor of shape (experts_per_tok, out_features), where `experts_per_tok` is the
|
|
number of activated experts per token.
|
|
"""
|
|
(local_experts, out_features, in_features), dtype = w.shape, w.dtype
|
|
_, experts_per_tok = indptr.shape
|
|
x_leading_dim, _ = x.shape
|
|
|
|
def access_x(x, e, j):
|
|
return x[0, j] if x_leading_dim == 1 else x[e, j]
|
|
|
|
# NOTE: Currently it assumes x.dtype == w.dtype, but the constraint can be relaxed easily.
|
|
assert w.shape == [local_experts, out_features, in_features] and w.dtype == dtype
|
|
assert x.shape == [x_leading_dim, in_features] and x.dtype == dtype
|
|
assert indptr.shape == [1, experts_per_tok] and indptr.dtype == "int32"
|
|
assert x_leading_dim in [1, experts_per_tok]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
x: T.Buffer((x_leading_dim, in_features), dtype),
|
|
w: T.Buffer((local_experts, out_features, in_features), dtype),
|
|
indptr: T.Buffer((1, experts_per_tok), "int32"),
|
|
o: T.Buffer((experts_per_tok, out_features), dtype),
|
|
):
|
|
T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable
|
|
for e in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
|
|
with T.sblock("gemv_o"):
|
|
e = T.axis.spatial(experts_per_tok, e)
|
|
T.reads(x[:, :], w[indptr[0, e], :, :], indptr[0, e])
|
|
T.writes(o[e, :])
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("gemv"):
|
|
i, j = T.axis.remap("SR", [i1, i2])
|
|
with T.init():
|
|
o[e, i] = T.cast(T.float16(0), dtype)
|
|
o[e, i] += access_x(x, e, j) * w[indptr[0, e], i, j]
|
|
|
|
return op.tensor_ir_op(
|
|
_func,
|
|
"moe_gemv",
|
|
args=[x, w, indptr],
|
|
out=Tensor.placeholder([experts_per_tok, out_features], dtype),
|
|
)
|
|
|
|
|
|
def dequantize_gemv(
|
|
x: Tensor,
|
|
w: Tensor,
|
|
scale: Tensor,
|
|
indptr: Tensor,
|
|
quantize_dtype: str,
|
|
group_size: int,
|
|
) -> Tensor:
|
|
"""GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized.
|
|
It needs to be dequantized before the GEMV computation.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
For project-in, the input tensor of shape (1, in_features); and for project-out, the input
|
|
shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
|
|
experts per token.
|
|
|
|
w : Tensor
|
|
The quantized weight tensor of shape (local_experts, out_features, in_features // n),
|
|
where n is the number of elements per storage dtype, e.g. if the storage dtype is uint32,
|
|
and the quantize dtype is int4, then n is 8.
|
|
`local_experts` is the total number of experts including activated and non-active ones.
|
|
|
|
scale : Tensor
|
|
The scale tensor of shape (local_experts, out_features, in_features // group_size), where
|
|
`local_experts` is the total number of experts including activated and non-active ones.
|
|
|
|
indptr : Tensor
|
|
The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
|
|
number of activated experts per token.
|
|
|
|
quantize_dtype : str
|
|
The quantize dtype of the weight tensor, which is usually int3, int4 or fp8, etc.
|
|
|
|
group_size : int
|
|
The number of elements in each quantization group, e.g. 32 or 128.
|
|
|
|
Returns
|
|
-------
|
|
out : Tensor
|
|
The output tensor of shape (experts_per_tok, out_features), where `experts_per_tok` is the
|
|
number of activated experts per token.
|
|
"""
|
|
(x_leading_dim, in_features), model_dtype = x.shape, x.dtype
|
|
(local_experts, out_features, _), storage_dtype = w.shape, w.dtype
|
|
_, experts_per_tok = indptr.shape
|
|
quantize_dtype_bits = DataType(quantize_dtype).bits
|
|
num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits
|
|
num_group = (in_features + group_size - 1) // group_size
|
|
num_storage = group_size // num_elem_per_storage * num_group
|
|
|
|
def _dequantize(w, s, e, i, j):
|
|
tir_bin_mask = tirx.const((2**quantize_dtype_bits) - 1, storage_dtype)
|
|
tir_max_int = tirx.const((2 ** (quantize_dtype_bits - 1)) - 1, model_dtype)
|
|
w = w[e, i, j // num_elem_per_storage]
|
|
s = s[e, i, j // group_size]
|
|
shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype)
|
|
w = tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype(model_dtype)
|
|
return (w - tir_max_int) * s
|
|
|
|
def access_x(x, e, j):
|
|
return x[0, j] if x_leading_dim == 1 else x[e, j]
|
|
|
|
assert x.shape == [x_leading_dim, in_features] and x.dtype == model_dtype
|
|
assert w.shape == [local_experts, out_features, num_storage] and w.dtype == storage_dtype
|
|
assert scale.shape == [local_experts, out_features, num_group] and scale.dtype == model_dtype
|
|
assert indptr.shape == [1, experts_per_tok] and indptr.dtype == "int32"
|
|
assert x_leading_dim in [1, experts_per_tok]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
x: T.Buffer((x_leading_dim, in_features), model_dtype),
|
|
w: T.Buffer((local_experts, out_features, num_storage), storage_dtype),
|
|
scale: T.Buffer((local_experts, out_features, num_group), model_dtype),
|
|
indptr: T.Buffer((1, experts_per_tok), "int32"),
|
|
o: T.Buffer((experts_per_tok, out_features), model_dtype),
|
|
):
|
|
T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable
|
|
for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
|
|
with T.sblock("gemv_o"):
|
|
e = T.axis.spatial(experts_per_tok, expert_id)
|
|
y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("dequantize"):
|
|
i, j = T.axis.remap("SS", [i1, i2])
|
|
y[i, j] = _dequantize(w, scale, indptr[0, e], i, j)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("gemv"):
|
|
i, j = T.axis.remap("SR", [i1, i2])
|
|
with T.init():
|
|
o[e, i] = T.cast(T.float16(0), model_dtype)
|
|
o[e, i] += access_x(x, e, j) * y[i, j]
|
|
|
|
return op.tensor_ir_op(
|
|
_func,
|
|
"moe_dequantize_gemv",
|
|
args=[x, w, scale, indptr],
|
|
out=Tensor.placeholder([experts_per_tok, out_features], model_dtype),
|
|
)
|
|
|
|
|
|
def dequantize_float8_gemv(
|
|
x: Tensor,
|
|
w: Tensor,
|
|
scale: Optional[Tensor],
|
|
indptr: Tensor,
|
|
quantize_dtype: Literal["float8_e5m2", "float8_e4m3fn"],
|
|
) -> Tensor:
|
|
"""GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized in
|
|
fp8 e5m2 or e4m3. It needs to be dequantized before the GEMV computation.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
For project-in, the input tensor of shape (1, in_features); and for project-out, the input
|
|
shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
|
|
experts per token.
|
|
|
|
w : Tensor
|
|
The quantized weight tensor of shape (local_experts, out_features, in_features)
|
|
|
|
scale : Optional[Tensor]
|
|
The optional scale tensor of shape (1,)
|
|
|
|
indptr : Tensor
|
|
The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
|
|
number of activated experts per token.
|
|
|
|
quantize_dtype : Literal["float8_e5m2", "float8_e4m3fn"]
|
|
The quantize dtype of the weight tensor, which is either float8_e5m2 or float8_e4m3fn.
|
|
"""
|
|
(x_leading_dim, in_features), model_dtype = x.shape, x.dtype
|
|
(local_experts, out_features, _), storage_dtype = w.shape, w.dtype
|
|
_, experts_per_tok = indptr.shape
|
|
quantize_dtype_bits = DataType(quantize_dtype).bits
|
|
num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits
|
|
num_storage = tirx.ceildiv(in_features, num_elem_per_storage)
|
|
|
|
def _dequantize(w, s, e, i, j):
|
|
if num_elem_per_storage == 1:
|
|
w = tirx.reinterpret(quantize_dtype, w[e, i, j])
|
|
else:
|
|
assert DataType(storage_dtype).type_code == DataTypeCode.UINT
|
|
tir_bin_mask = tirx.const((2**quantize_dtype_bits) - 1, storage_dtype)
|
|
w = w[e, i, j // num_elem_per_storage]
|
|
shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype)
|
|
w = tirx.reinterpret(
|
|
quantize_dtype,
|
|
tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype("uint8"),
|
|
)
|
|
w = w.astype(model_dtype)
|
|
if s is not None:
|
|
w = w * s[0]
|
|
return w
|
|
|
|
def access_x(x, e, j):
|
|
return x[0, j] if x_leading_dim == 1 else x[e, j]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func_with_scale(
|
|
x: T.Buffer((x_leading_dim, in_features), model_dtype),
|
|
w: T.Buffer((local_experts, out_features, num_storage), storage_dtype),
|
|
scale: T.Buffer((1,), "float32"),
|
|
indptr: T.Buffer((1, experts_per_tok), "int32"),
|
|
o: T.Buffer((experts_per_tok, out_features), model_dtype),
|
|
):
|
|
T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable
|
|
for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
|
|
with T.sblock("gemv_o"):
|
|
e = T.axis.spatial(experts_per_tok, expert_id)
|
|
y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("dequantize"):
|
|
i, j = T.axis.remap("SS", [i1, i2])
|
|
y[i, j] = _dequantize(w, scale, indptr[0, e], i, j)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("gemv"):
|
|
i, j = T.axis.remap("SR", [i1, i2])
|
|
with T.init():
|
|
o[e, i] = T.cast(T.float16(0), model_dtype)
|
|
o[e, i] += access_x(x, e, j) * y[i, j]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func_without_scale(
|
|
x: T.Buffer((x_leading_dim, in_features), model_dtype),
|
|
w: T.Buffer((local_experts, out_features, num_storage), storage_dtype),
|
|
indptr: T.Buffer((1, experts_per_tok), "int32"),
|
|
o: T.Buffer((experts_per_tok, out_features), model_dtype),
|
|
):
|
|
T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable
|
|
for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
|
|
with T.sblock("gemv_o"):
|
|
e = T.axis.spatial(experts_per_tok, expert_id)
|
|
y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("dequantize"):
|
|
i, j = T.axis.remap("SS", [i1, i2])
|
|
y[i, j] = _dequantize(w, None, indptr[0, e], i, j)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("gemv"):
|
|
i, j = T.axis.remap("SR", [i1, i2])
|
|
with T.init():
|
|
o[e, i] = T.cast(T.float16(0), model_dtype)
|
|
o[e, i] += access_x(x, e, j) * y[i, j]
|
|
|
|
if scale is not None:
|
|
return op.tensor_ir_op(
|
|
_func_with_scale,
|
|
"moe_dequantize_gemv",
|
|
args=[x, w, scale, indptr],
|
|
out=Tensor.placeholder([experts_per_tok, out_features], model_dtype),
|
|
)
|
|
return op.tensor_ir_op(
|
|
_func_without_scale,
|
|
"moe_dequantize_gemv",
|
|
args=[x, w, indptr],
|
|
out=Tensor.placeholder([experts_per_tok, out_features], model_dtype),
|
|
)
|
|
|
|
|
|
def dequantize_block_scale_float8_gemv(
|
|
x: Tensor,
|
|
w: Tensor,
|
|
w_scale: Tensor,
|
|
expert_indices: Tensor,
|
|
block_size: Tuple[int, int], # noqa: UP006
|
|
out_dtype: str,
|
|
) -> Tensor:
|
|
"""GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized in
|
|
fp8 e5m2 or e4m3. It needs to be dequantized before the GEMV computation.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
For project-in, the input tensor of shape (1, in_features); and for project-out, the input
|
|
shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
|
|
experts per token.
|
|
|
|
w : Tensor
|
|
The quantized weight tensor of shape (local_experts, out_features, in_features)
|
|
|
|
w_scale : Tensor
|
|
The scale tensor of shape
|
|
(local_experts, out_features // block_size[0], in_features // block_size[1])
|
|
|
|
indptr : Tensor
|
|
The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
|
|
number of activated experts per token.
|
|
|
|
block_size : Tuple[int, int]
|
|
The block size of the weight tensor.
|
|
|
|
out_dtype : str
|
|
The output dtype of the GEMV computation.
|
|
"""
|
|
x_leading_dim, in_features = x.shape
|
|
local_experts, out_features, k = w.shape
|
|
_, experts_per_tok = expert_indices.shape
|
|
model_dtype = x.dtype
|
|
quantize_dtype = w.dtype
|
|
|
|
assert out_features % block_size[0] == 0
|
|
assert k % block_size[1] == 0
|
|
|
|
def _dequantize(w, s, e, i, j):
|
|
return w[e, i, j].astype(model_dtype) * s[e, i // block_size[0], j // block_size[1]].astype(
|
|
model_dtype
|
|
)
|
|
|
|
def load_x(x, e, j):
|
|
return x[0, j] if x_leading_dim == 1 else x[e, j]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
x: T.Buffer((x_leading_dim, in_features), model_dtype),
|
|
w: T.Buffer((local_experts, out_features, k), quantize_dtype),
|
|
w_scale: T.Buffer(
|
|
(local_experts, out_features // block_size[0], k // block_size[1]),
|
|
"float32",
|
|
),
|
|
expert_indices: T.Buffer((1, experts_per_tok), "int32"),
|
|
o: T.Buffer((experts_per_tok, out_features), out_dtype),
|
|
):
|
|
T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable
|
|
for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
|
|
with T.sblock("gemv_o"):
|
|
e = T.axis.spatial(experts_per_tok, expert_id)
|
|
y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("dequantize"):
|
|
i, j = T.axis.remap("SS", [i1, i2])
|
|
y[i, j] = _dequantize(w, w_scale, expert_indices[0, e], i, j)
|
|
for i1, i2 in T.grid(out_features, in_features):
|
|
with T.sblock("gemv"):
|
|
i, j = T.axis.remap("SR", [i1, i2])
|
|
with T.init():
|
|
o[e, i] = T.cast(T.float16(0), out_dtype)
|
|
o[e, i] += (load_x(x, e, j) * y[i, j]).astype(out_dtype)
|
|
|
|
return op.tensor_ir_op(
|
|
_func,
|
|
"moe_dequantize_gemv",
|
|
args=[x, w, w_scale, expert_indices],
|
|
out=Tensor.placeholder([experts_per_tok, out_features], out_dtype),
|
|
)
|
|
|
|
|
|
def group_gemm(x: Tensor, w: Tensor, indptr: Tensor):
|
|
"""Group GEMM in MoE models.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
Input tensor of shape (batch_size, in_features), where `batch_size` could be dynamic shape.
|
|
|
|
w : Tensor
|
|
Weight tensor of shape (num_local_experts, out_features, in_features).
|
|
`w[i, :, :]` is the weight matrix for the `i`-th local expert.
|
|
|
|
indptr : Tensor
|
|
Index pointer tensor of shape (num_local_experts + 1, ).
|
|
`x[indptr[a] : indptr[a + 1]]` is the input for the `i`-th local expert.
|
|
|
|
Returns
|
|
-------
|
|
out : Tensor
|
|
Output tensor of shape (batch_size, out_features).
|
|
"""
|
|
# NOTE: Currently it assumes x.dtype == w.dtype, but the constraint can be relaxed easily.
|
|
(num_local_experts, out_features, in_features), dtype = w.shape, w.dtype
|
|
|
|
assert x.shape[1:] == [in_features] and x.dtype == dtype
|
|
assert indptr.shape == [num_local_experts + 1] and indptr.dtype == "int32"
|
|
|
|
Ne, N, K = num_local_experts, out_features, in_features
|
|
BLK_M, BLK_N, BLK_K = 8, 128, 32
|
|
TX, TY, CTA_COUNT = 8, 32, 1024
|
|
VEC_X, VEC_W, VEC_O, VEC_DOT = 1, 1, 1, 1
|
|
UNROLL = 64
|
|
STORAGE_ALIGN = False
|
|
assert BLK_K % 8 == 0
|
|
tiles_per_row = (N + BLK_N - 1) // BLK_N
|
|
zero = tirx.const(0, dtype)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
var_x: T.handle,
|
|
var_w: T.handle,
|
|
var_indptr: T.handle,
|
|
var_o: T.handle,
|
|
):
|
|
T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
|
|
B = T.int32()
|
|
X = T.match_buffer(var_x, (B, K), dtype)
|
|
W = T.match_buffer(var_w, (Ne, N, K), dtype)
|
|
indptr = T.match_buffer(var_indptr, (Ne + 1,), "int32")
|
|
out = T.match_buffer(var_o, (B, N), dtype)
|
|
|
|
for _bx in T.thread_binding(CTA_COUNT, thread="blockIdx.x"):
|
|
with T.sblock("CTA"):
|
|
bx = T.axis.spatial(CTA_COUNT, _bx)
|
|
T.reads(indptr[:], X[:, :], W[:, :, :])
|
|
T.writes(out[:, :])
|
|
sum = T.sblock_alloc_buffer((2,), "int32", scope="local")
|
|
row = T.sblock_alloc_buffer((2,), "int32", scope="local")
|
|
cur_e = T.sblock_alloc_buffer((1,), "int32", scope="local")
|
|
tile_id = T.sblock_alloc_buffer((1,), "int32", scope="local")
|
|
sum[0] = 0
|
|
sum[1] = T.ceildiv(indptr[1] - indptr[0], BLK_M) * tiles_per_row
|
|
row[0] = 0
|
|
row[1] = indptr[1] - indptr[0]
|
|
cur_e[0] = 0
|
|
tile_id[0] = bx
|
|
while T.tvm_thread_invariant(cur_e[0] < Ne):
|
|
# move to the current group
|
|
while sum[1] <= tile_id[0] and cur_e[0] < Ne:
|
|
cur_e[0] += 1
|
|
if cur_e[0] < Ne:
|
|
e: T.int32 = cur_e[0]
|
|
delta: T.int32 = indptr[e + 1] - indptr[e]
|
|
sum[0] = sum[1]
|
|
sum[1] += T.ceildiv(delta, BLK_M) * tiles_per_row
|
|
row[0] = row[1]
|
|
row[1] += delta
|
|
# sync threads to make sure all threads have the same tile position
|
|
T.tvm_storage_sync("shared")
|
|
if T.tvm_thread_invariant(cur_e[0] < Ne):
|
|
# fetch current tile position
|
|
e: T.int32 = cur_e[0]
|
|
num_tiles: T.int32 = tile_id[0] - sum[0]
|
|
m_offset: T.int32 = BLK_M * T.floordiv(num_tiles, tiles_per_row) + row[0]
|
|
n_offset: T.int32 = BLK_N * T.floormod(num_tiles, tiles_per_row)
|
|
with T.sblock("gemm"):
|
|
T.reads(
|
|
row[1],
|
|
X[m_offset : m_offset + BLK_M, :],
|
|
W[e, n_offset : n_offset + BLK_N, :],
|
|
)
|
|
T.writes(
|
|
out[
|
|
m_offset : m_offset + BLK_M,
|
|
n_offset : n_offset + BLK_N,
|
|
]
|
|
)
|
|
X_tile = T.sblock_alloc_buffer((BLK_M, K), dtype, scope="shared")
|
|
W_tile = T.sblock_alloc_buffer((BLK_N, K), dtype, scope="shared")
|
|
O_tile = T.sblock_alloc_buffer((BLK_M, BLK_N), dtype, scope="local")
|
|
for a0, a1 in T.grid(BLK_M, K):
|
|
with T.sblock("X_shared"):
|
|
i, j = T.axis.remap("SS", [a0, a1])
|
|
X_tile[i, j] = T.if_then_else(
|
|
m_offset + i < row[1],
|
|
X[m_offset + i, j],
|
|
zero,
|
|
)
|
|
for a0, a1 in T.grid(BLK_N, K):
|
|
with T.sblock("W_shared"):
|
|
i, j = T.axis.remap("SS", [a0, a1])
|
|
W_tile[i, j] = T.if_then_else(
|
|
n_offset + i < N,
|
|
W[e, n_offset + i, j],
|
|
zero,
|
|
)
|
|
for a0, a1, a2 in T.grid(BLK_M, BLK_N, K):
|
|
with T.sblock("compute"):
|
|
i, j, k = T.axis.remap("SSR", [a0, a1, a2])
|
|
with T.init():
|
|
O_tile[i, j] = zero
|
|
O_tile[i, j] += X_tile[i, k] * W_tile[j, k]
|
|
for a0, a1 in T.grid(BLK_M, BLK_N):
|
|
with T.sblock("store"):
|
|
i, j = T.axis.remap("SS", [a0, a1])
|
|
if m_offset + i < row[1] and n_offset + j < N:
|
|
out[m_offset + i, n_offset + j] = O_tile[i, j]
|
|
# move to next tile
|
|
tile_id[0] += CTA_COUNT
|
|
|
|
def _schedule():
|
|
sch = s_tir.Schedule(_func)
|
|
|
|
def _cooperative_fetch(block, vec_len):
|
|
num_loops = len(sch.get_loops(block))
|
|
sch.compute_at(block, ko, preserve_unit_loops=True)
|
|
loops = sch.get_loops(block)[-num_loops:]
|
|
ty, tx, _, vec = sch.split(
|
|
sch.fuse(*loops),
|
|
factors=[TY, TX, None, vec_len],
|
|
)
|
|
sch.vectorize(vec)
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
if STORAGE_ALIGN:
|
|
sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len)
|
|
return block
|
|
|
|
main_block = sch.get_sblock("compute")
|
|
x, y, k = sch.get_loops(main_block)
|
|
ty, yi = sch.split(y, [TY, None])
|
|
tx, xi, vec_c = sch.split(x, [TX, None, VEC_DOT])
|
|
ko, ki = sch.split(k, factors=[None, BLK_K])
|
|
sch.reorder(ty, tx, ko, ki, yi, xi, vec_c)
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
sch.vectorize(vec_c)
|
|
if UNROLL > 0:
|
|
sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=UNROLL)
|
|
sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
|
|
l2g = sch.get_sblock("store")
|
|
sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True)
|
|
_, v = sch.split(sch.get_loops(l2g)[-1], [None, VEC_O])
|
|
sch.vectorize(v)
|
|
_cooperative_fetch(sch.get_sblock("X_shared"), vec_len=VEC_X)
|
|
_cooperative_fetch(sch.get_sblock("W_shared"), vec_len=VEC_W)
|
|
sch.decompose_reduction(main_block, ko)
|
|
return sch.mod["main"]
|
|
|
|
return op.tensor_ir_op(
|
|
_schedule(),
|
|
"group_gemm",
|
|
args=[x, w, indptr],
|
|
out=Tensor.placeholder([x.shape[0], out_features], dtype),
|
|
)
|
|
|
|
|
|
def dequantize_group_gemm(
|
|
x: Tensor,
|
|
w: Tensor,
|
|
scale: Tensor,
|
|
indptr: Tensor,
|
|
quantize_dtype: str,
|
|
indptr_dtype: str,
|
|
group_size: int,
|
|
):
|
|
"""Group GEMM in MoE models but the weight is quantized.
|
|
|
|
Parameters
|
|
----------
|
|
x : Tensor
|
|
Input tensor of shape (batch_size, in_features), where `batch_size` could be dynamic shape.
|
|
|
|
w : Tensor
|
|
Weight tensor of shape (num_local_experts, out_features, in_features // n), where n is the
|
|
number of elements per storage dtype, e.g. if the storage dtype is uint32, and the quantize
|
|
dtype is int4, then n is 8.
|
|
|
|
scale : Tensor
|
|
The scale tensor of shape (num_local_experts, out_features, in_features // group_size).
|
|
|
|
indptr : Tensor
|
|
Index pointer tensor of shape (num_local_experts + 1, ). `x[indptr[a] : indptr[a + 1]]` is
|
|
the input for the `i`-th local expert.
|
|
|
|
group_size : int
|
|
The number of elements in each quantization group, e.g. 32 or 128.
|
|
|
|
quantize_dtype : str
|
|
The quantize dtype of the weight tensor, which is usually int3, int4 or fp8, etc.
|
|
|
|
indptr_dtype : str
|
|
The dtype of the index pointer tensor, which can be int32 or int64.
|
|
|
|
Returns
|
|
-------
|
|
out : Tensor
|
|
Output tensor of shape (batch_size, out_features).
|
|
"""
|
|
(_, in_features), model_dtype = x.shape, x.dtype
|
|
(num_local_experts, out_features, _), storage_dtype = w.shape, w.dtype
|
|
quantize_dtype_bits = DataType(quantize_dtype).bits
|
|
num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits
|
|
num_group = (in_features + group_size - 1) // group_size
|
|
num_storage = group_size // num_elem_per_storage * num_group
|
|
|
|
def _dequantize(w, s, e, i, j):
|
|
tir_bin_mask = tirx.const((1 << quantize_dtype_bits) - 1, storage_dtype)
|
|
tir_max_int = tirx.const((2 ** (quantize_dtype_bits - 1)) - 1, model_dtype)
|
|
w = w[e, i, j // num_elem_per_storage]
|
|
s = s[e, i, j // group_size]
|
|
shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype)
|
|
w = tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype(model_dtype)
|
|
return (w - tir_max_int) * s
|
|
|
|
Ne, N, K = num_local_experts, out_features, in_features
|
|
BLK_M, BLK_N, BLK_K = 8, 128, 32
|
|
TX, TY, CTA_COUNT = 8, 32, 1024
|
|
VEC_X, VEC_W, VEC_O, VEC_DOT = 1, 1, 1, 1
|
|
UNROLL = 64
|
|
STORAGE_ALIGN = False
|
|
assert BLK_K % 8 == 0
|
|
tiles_per_row = (N + BLK_N - 1) // BLK_N
|
|
zero = tirx.const(0, model_dtype)
|
|
if indptr_dtype == "int64":
|
|
indptr = op.pad(indptr, [1, 0], "constant", 0)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _func(
|
|
var_x: T.handle,
|
|
w: T.Buffer((Ne, N, num_storage), storage_dtype),
|
|
scale: T.Buffer((Ne, N, num_group), model_dtype),
|
|
indptr: T.Buffer((Ne + 1,), indptr_dtype),
|
|
var_o: T.handle,
|
|
):
|
|
T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
|
|
B = T.int32()
|
|
X = T.match_buffer(var_x, (B, K), model_dtype)
|
|
out = T.match_buffer(var_o, (B, N), model_dtype)
|
|
for _bx in T.thread_binding(CTA_COUNT, thread="blockIdx.x"):
|
|
with T.sblock("CTA"):
|
|
bx = T.axis.spatial(CTA_COUNT, _bx)
|
|
T.reads(X[:, :], w[:, :, :], scale[:, :, :], indptr[:])
|
|
T.writes(out[:, :])
|
|
sum = T.sblock_alloc_buffer((2,), indptr_dtype, scope="local")
|
|
row = T.sblock_alloc_buffer((2,), indptr_dtype, scope="local")
|
|
cur_e = T.sblock_alloc_buffer((1,), indptr_dtype, scope="local")
|
|
tile_id = T.sblock_alloc_buffer((1,), indptr_dtype, scope="local")
|
|
sum[0] = 0
|
|
sum[1] = T.ceildiv(indptr[1] - indptr[0], BLK_M) * tiles_per_row
|
|
row[0] = 0
|
|
row[1] = indptr[1] - indptr[0]
|
|
cur_e[0] = 0
|
|
tile_id[0] = bx
|
|
while T.tvm_thread_invariant(cur_e[0] < Ne):
|
|
# move to the current group
|
|
while sum[1] <= tile_id[0] and cur_e[0] < Ne:
|
|
cur_e[0] += 1
|
|
if cur_e[0] < Ne:
|
|
e = cur_e[0]
|
|
delta = indptr[e + 1] - indptr[e]
|
|
sum[0] = sum[1]
|
|
sum[1] += T.ceildiv(delta, BLK_M) * tiles_per_row
|
|
row[0] = row[1]
|
|
row[1] += delta
|
|
# sync threads to make sure all threads have the same tile position
|
|
T.tvm_storage_sync("shared")
|
|
if T.tvm_thread_invariant(cur_e[0] < Ne):
|
|
# fetch current tile position
|
|
e = cur_e[0]
|
|
num_tiles = tile_id[0] - sum[0]
|
|
m_offset = T.floordiv(num_tiles, tiles_per_row) * BLK_M + row[0]
|
|
n_offset = T.floormod(num_tiles, tiles_per_row) * BLK_N
|
|
with T.sblock("gemm"):
|
|
T.reads(
|
|
row[1],
|
|
X[m_offset : m_offset + BLK_M, :],
|
|
w[e, n_offset : n_offset + BLK_N, :],
|
|
scale[e, n_offset : n_offset + BLK_N, :],
|
|
)
|
|
T.writes(
|
|
out[
|
|
m_offset : m_offset + BLK_M,
|
|
n_offset : n_offset + BLK_N,
|
|
]
|
|
)
|
|
X_tile = T.sblock_alloc_buffer((BLK_M, K), model_dtype, scope="shared")
|
|
W_tile = T.sblock_alloc_buffer((BLK_N, K), model_dtype, scope="shared")
|
|
O_tile = T.sblock_alloc_buffer((BLK_M, BLK_N), "float32", scope="local")
|
|
for a0, a1 in T.grid(BLK_M, K):
|
|
with T.sblock("X_shared"):
|
|
i, j = T.axis.remap("SS", [a0, a1])
|
|
X_tile[i, j] = T.if_then_else(
|
|
m_offset + i < row[1],
|
|
X[m_offset + i, j],
|
|
zero,
|
|
)
|
|
for a0, a1 in T.grid(BLK_N, K):
|
|
with T.sblock("W_shared"):
|
|
i, j = T.axis.remap("SS", [a0, a1])
|
|
W_tile[i, j] = T.if_then_else(
|
|
n_offset + i < N,
|
|
_dequantize(w, scale, e, n_offset + i, j),
|
|
zero,
|
|
)
|
|
for a0, a1, a2 in T.grid(BLK_M, BLK_N, K):
|
|
with T.sblock("compute"):
|
|
i, j, k = T.axis.remap("SSR", [a0, a1, a2])
|
|
with T.init():
|
|
O_tile[i, j] = zero
|
|
O_tile[i, j] += X_tile[i, k] * W_tile[j, k]
|
|
for a0, a1 in T.grid(BLK_M, BLK_N):
|
|
with T.sblock("store"):
|
|
i, j = T.axis.remap("SS", [a0, a1])
|
|
if m_offset + i < row[1] and n_offset + j < N:
|
|
out[m_offset + i, n_offset + j] = O_tile[i, j]
|
|
# move to next tile
|
|
tile_id[0] += CTA_COUNT
|
|
|
|
def _schedule():
|
|
sch = s_tir.Schedule(_func)
|
|
|
|
def _cooperative_fetch(block, vec_len):
|
|
num_loops = len(sch.get_loops(block))
|
|
sch.compute_at(block, ko, preserve_unit_loops=True)
|
|
loops = sch.get_loops(block)[-num_loops:]
|
|
ty, tx, _, vec = sch.split(
|
|
sch.fuse(*loops),
|
|
factors=[TY, TX, None, vec_len],
|
|
)
|
|
sch.vectorize(vec)
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
if STORAGE_ALIGN:
|
|
sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len)
|
|
return block
|
|
|
|
main_block = sch.get_sblock("compute")
|
|
x, y, k = sch.get_loops(main_block)
|
|
ty, yi = sch.split(y, [TY, None])
|
|
tx, xi, vec_c = sch.split(x, [TX, None, VEC_DOT])
|
|
ko, ki = sch.split(k, factors=[None, BLK_K])
|
|
sch.reorder(ty, tx, ko, ki, yi, xi, vec_c)
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
sch.vectorize(vec_c)
|
|
if UNROLL > 0:
|
|
sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=UNROLL)
|
|
sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
|
|
l2g = sch.get_sblock("store")
|
|
sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True)
|
|
_, v = sch.split(sch.get_loops(l2g)[-1], [None, VEC_O])
|
|
sch.vectorize(v)
|
|
_cooperative_fetch(sch.get_sblock("X_shared"), vec_len=VEC_X)
|
|
_cooperative_fetch(sch.get_sblock("W_shared"), vec_len=VEC_W)
|
|
sch.decompose_reduction(main_block, ko)
|
|
return sch.mod["main"]
|
|
|
|
return op.tensor_ir_op(
|
|
_schedule(),
|
|
"dequantize_group_gemm",
|
|
args=[x, w, scale, indptr],
|
|
out=Tensor.placeholder([x.shape[0], out_features], model_dtype),
|
|
)
|