chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,4 @@
|
||||
"""Common `nn.Modules` used to define LLMs in this project."""
|
||||
|
||||
from .expert import MixtralExperts
|
||||
from .kv_cache import PagedKVCache, RopeMode
|
||||
@@ -0,0 +1,33 @@
|
||||
"""An nn.Module that represents MoE experts"""
|
||||
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.relax.frontend.nn import Tensor
|
||||
|
||||
from mlc_llm.op import cutlass, extern, ft_gemm, moe_matmul
|
||||
|
||||
|
||||
class MixtralExperts(nn.Module):
|
||||
"""Mixtral experts"""
|
||||
|
||||
def __init__(self, num_local_experts, in_features, out_features, tensor_parallel_shards=1):
|
||||
self.num_local_experts = num_local_experts
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.weight = nn.Parameter((num_local_experts, out_features, in_features))
|
||||
self.dtype = "float32"
|
||||
self.tensor_parallel_shards = tensor_parallel_shards
|
||||
|
||||
def forward(self, x: Tensor, indptr: Tensor):
|
||||
assert x.ndim == 2
|
||||
if indptr.ndim == 2:
|
||||
assert indptr.shape[0] == 1
|
||||
return moe_matmul.gemv(x, self.weight, indptr)
|
||||
assert indptr.ndim == 1
|
||||
if extern.get_store().cutlass_group_gemm and self.dtype in [
|
||||
"float16",
|
||||
"bfloat16",
|
||||
]:
|
||||
return cutlass.group_gemm(x, self.weight, indptr)
|
||||
if extern.get_store().faster_transformer and self.dtype == "float16":
|
||||
return ft_gemm.faster_transformer_moe_gemm(x, self.weight, indptr)
|
||||
return moe_matmul.group_gemm(x, self.weight, indptr)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Attention KV cache modeling."""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Literal, Optional, Union # noqa: UP035
|
||||
|
||||
import numpy as np
|
||||
from tvm import relax as rx
|
||||
from tvm import tirx
|
||||
from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache as TVMPagedKVCache
|
||||
from tvm.relax.frontend.nn.llm.kv_cache import RopeMode
|
||||
|
||||
|
||||
class PagedKVCache(TVMPagedKVCache):
|
||||
"""The Paged KV Cache used in LLM batching for efficient attention computation."""
|
||||
|
||||
@staticmethod
|
||||
def create_generic(
|
||||
attn_kind: Union[Literal["mha", "mla"], List[Literal["mha", "mla", "mha_sliding"]]], # noqa: UP006
|
||||
max_batch_size: tirx.Var,
|
||||
max_total_seq_len: tirx.Var,
|
||||
prefill_chunk_size: tirx.Var,
|
||||
page_size: tirx.Var,
|
||||
support_sliding_window: tirx.Var,
|
||||
num_hidden_layers: int,
|
||||
num_attention_heads: int,
|
||||
num_key_value_heads: int,
|
||||
qk_head_dim: int,
|
||||
v_head_dim: int,
|
||||
rope_mode: RopeMode,
|
||||
rope_scale: int,
|
||||
rope_theta: int,
|
||||
dtype: str,
|
||||
mla_original_qk_head_dim: int = 0,
|
||||
mla_original_v_head_dim: int = 0,
|
||||
rotary_dim: Optional[int] = None,
|
||||
rope_scaling: Optional[Dict[str, Any]] = None, # noqa: UP006
|
||||
rope_ext_factors: Optional[List[int]] = None, # noqa: UP006
|
||||
layer_partition: Optional[List[int]] = None, # noqa: UP006
|
||||
enable_disaggregation: bool = False,
|
||||
name: str = "paged_kv_cache",
|
||||
) -> "PagedKVCache":
|
||||
"""The generic function of creating a multi-head attention PagedKVCache,
|
||||
which will be rewritten by functions in compilation pipeline.
|
||||
"""
|
||||
if rotary_dim is None:
|
||||
rotary_dim = qk_head_dim
|
||||
if rope_scaling is None:
|
||||
rope_scaling = {}
|
||||
if layer_partition is None:
|
||||
layer_partition = [0, num_hidden_layers]
|
||||
if isinstance(attn_kind, List): # noqa: UP006
|
||||
rx_attn_kind = [rx.StringImm(layer_kind) for layer_kind in attn_kind]
|
||||
else:
|
||||
rx_attn_kind = rx.StringImm(attn_kind)
|
||||
return PagedKVCache(
|
||||
_expr=rx.call_pure_packed(
|
||||
"mlc.create_paged_kv_cache_generic",
|
||||
rx_attn_kind,
|
||||
rx.ShapeExpr(
|
||||
[
|
||||
max_batch_size,
|
||||
max_total_seq_len,
|
||||
prefill_chunk_size,
|
||||
page_size,
|
||||
support_sliding_window,
|
||||
]
|
||||
),
|
||||
rx.ShapeExpr(layer_partition),
|
||||
rx.prim_value(num_hidden_layers),
|
||||
rx.prim_value(num_attention_heads),
|
||||
rx.prim_value(num_key_value_heads),
|
||||
rx.prim_value(qk_head_dim),
|
||||
rx.prim_value(v_head_dim),
|
||||
rx.prim_value(mla_original_qk_head_dim),
|
||||
rx.prim_value(mla_original_v_head_dim),
|
||||
rx.prim_value(rope_mode),
|
||||
rx.prim_value(rope_scale),
|
||||
rx.prim_value(rope_theta),
|
||||
rx.StringImm(json.dumps(rope_scaling)),
|
||||
(
|
||||
rx.const(np.array(rope_ext_factors, "float32"))
|
||||
if rope_ext_factors is not None
|
||||
else rx.prim_value(0)
|
||||
# NOTE: since relax does not have "Optional" type, we use prim_value(0)
|
||||
# to represent "undefined".
|
||||
),
|
||||
rx.prim_value(rotary_dim),
|
||||
rx.prim_value(int(enable_disaggregation)),
|
||||
rx.DataTypeImm(dtype),
|
||||
ty_args=rx.ObjectType(),
|
||||
),
|
||||
_name=name,
|
||||
)
|
||||
@@ -0,0 +1,337 @@
|
||||
"""RNN State modeling."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Union
|
||||
|
||||
from tvm import relax as rx
|
||||
from tvm import tirx
|
||||
from tvm.relax.frontend.nn import Object, Tensor
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
class RNNState(Object):
|
||||
"""The RNN State used in Space State Models"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
max_batch_size: tirx.Var,
|
||||
num_hidden_layers: int,
|
||||
max_history: int,
|
||||
init_values: Sequence[rx.Constant],
|
||||
name: str = "rnn_state",
|
||||
) -> "RNNState":
|
||||
"""Create a RNN state object.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
max_batch_size : tirx.Var
|
||||
The maximum batch size.
|
||||
num_hidden_layers : int
|
||||
The number of hidden layers.
|
||||
max_history : int
|
||||
The maximum history length.
|
||||
init_values : Sequence[rx.Constant]
|
||||
The initial values of the RNN state. Must be compile-time Relax constants
|
||||
(e.g. R.const(np.zeros(...))).
|
||||
"""
|
||||
|
||||
bb = rx.BlockBuilder.current()
|
||||
state_infos = [
|
||||
(tuple(int(x) for x in v.data.shape), str(v.data.dtype)) for v in init_values
|
||||
]
|
||||
|
||||
f_gets = [
|
||||
bb.add_func(
|
||||
RNNState.create_get_func(shape, dtype, max_batch_size, max_history, id),
|
||||
f"rnn_state_get_{id}",
|
||||
)
|
||||
for id, (shape, dtype) in enumerate(state_infos)
|
||||
]
|
||||
f_sets = [
|
||||
bb.add_func(
|
||||
RNNState.create_set_func(shape, dtype, max_batch_size, max_history, id),
|
||||
f"rnn_state_set_{id}",
|
||||
)
|
||||
for id, (shape, dtype) in enumerate(state_infos)
|
||||
]
|
||||
|
||||
ret = RNNState(
|
||||
_expr=rx.call_pure_packed(
|
||||
"vm.builtin.rnn_state_create",
|
||||
rx.prim_value(num_hidden_layers),
|
||||
max_batch_size,
|
||||
max_history,
|
||||
f_gets,
|
||||
f_sets,
|
||||
list(init_values),
|
||||
ty_args=[rx.ObjectType()],
|
||||
),
|
||||
_name=name,
|
||||
)
|
||||
return ret
|
||||
|
||||
def get(
|
||||
self,
|
||||
layer_id: int,
|
||||
state_id: int,
|
||||
shape: Sequence[tirx.Expr],
|
||||
dtype: str,
|
||||
) -> Tensor:
|
||||
"""Get the state of the RNN layer.
|
||||
|
||||
- If there is only one sequence, we can directly use the storage memory,
|
||||
without copying the data.
|
||||
- If there are multiple sequences, we need to copy the data to get a contiguous
|
||||
memory.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
layer_id : int
|
||||
The layer id.
|
||||
state_id : int
|
||||
The state id.
|
||||
shape : Sequence[tirx.Expr]
|
||||
The shape of the state tensor.
|
||||
dtype: str
|
||||
The data type of the state tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor
|
||||
The state tensor, with shape `(batch_size, *state_size)`.
|
||||
"""
|
||||
bb = rx.BlockBuilder.current()
|
||||
|
||||
return Tensor(
|
||||
_expr=bb.emit(
|
||||
rx.call_dps_packed(
|
||||
"vm.builtin.rnn_state_get",
|
||||
[self._expr, layer_id, state_id],
|
||||
out_ty=rx.TensorType(shape, dtype),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def set(self, layer_id: int, state_id: int, value: Tensor) -> "RNNState":
|
||||
"""Set the state of the RNN layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
layer_id : int
|
||||
The layer id.
|
||||
state_id : int
|
||||
The state id.
|
||||
value : Tensor
|
||||
The state tensor, with shape `(batch_size, *state_size)`.
|
||||
"""
|
||||
bb = rx.BlockBuilder.current()
|
||||
return RNNState(
|
||||
_expr=bb.emit(
|
||||
rx.call_pure_packed(
|
||||
"vm.builtin.rnn_state_set",
|
||||
self._expr,
|
||||
rx.prim_value(layer_id),
|
||||
rx.prim_value(state_id),
|
||||
value._expr,
|
||||
ty_args=[rx.ObjectType()],
|
||||
)
|
||||
),
|
||||
_name="rnn_state_set",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_get_func(
|
||||
shape: Sequence[Union[int, tirx.Var]],
|
||||
dtype: str,
|
||||
max_batch_size: Union[int, tirx.Var],
|
||||
max_history: Union[int, tirx.Var],
|
||||
state_id: int,
|
||||
) -> tirx.PrimFunc:
|
||||
"""Create the get function with given state shape.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape : Sequence[Union[int, tirx.Var]]
|
||||
The shape of the state tensor.
|
||||
|
||||
dtype: str
|
||||
The data type of the state tensor.
|
||||
|
||||
max_batch_size : Union[int, tirx.Var]
|
||||
The maximum batch size.
|
||||
|
||||
max_history : Union[int, tirx.Var]
|
||||
The maximum history length.
|
||||
|
||||
state_id : int
|
||||
The id of the state, used for naming the function.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tirx.PrimFunc
|
||||
The get function.
|
||||
"""
|
||||
|
||||
def _func_one_dim():
|
||||
@T.prim_func(s_tir=True)
|
||||
def f(
|
||||
var_storage: T.handle,
|
||||
var_seq_slot_ids: T.handle,
|
||||
var_history_slot_ids: T.handle,
|
||||
var_output: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
T.func_attr({"global_symbol": f"rnn_state_get_{state_id}"})
|
||||
|
||||
storage = T.match_buffer(
|
||||
var_storage, (max_batch_size, max_history, shape[0]), dtype
|
||||
)
|
||||
seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32")
|
||||
history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32")
|
||||
output = T.match_buffer(var_output, (batch_size, shape[0]), dtype)
|
||||
|
||||
for i in range(batch_size):
|
||||
for s in range(shape[0]):
|
||||
with T.sblock("copy"):
|
||||
vi, vs = T.axis.remap("SS", [i, s])
|
||||
seq_id: T.int32 = seq_slot_ids[vi]
|
||||
history_id: T.int32 = history_slot_ids[vi]
|
||||
output[vi, vs] = storage[seq_id, history_id, vs]
|
||||
|
||||
return f
|
||||
|
||||
def _func_high_dim():
|
||||
# Add a wrapper function to avoid parse the following code when len(shape) = 1
|
||||
@T.prim_func(s_tir=True)
|
||||
def f(
|
||||
var_storage: T.handle,
|
||||
var_seq_slot_ids: T.handle,
|
||||
var_history_slot_ids: T.handle,
|
||||
var_output: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
T.func_attr({"global_symbol": f"rnn_state_get_{state_id}"})
|
||||
|
||||
storage = T.match_buffer(var_storage, (max_batch_size, max_history, *shape), dtype)
|
||||
seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32")
|
||||
history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32")
|
||||
output = T.match_buffer(var_output, (batch_size, *shape), dtype)
|
||||
|
||||
for i in range(batch_size):
|
||||
for s in T.grid(*shape):
|
||||
with T.sblock("copy"):
|
||||
vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s])
|
||||
seq_id: T.int32 = seq_slot_ids[vi]
|
||||
history_id: T.int32 = history_slot_ids[vi]
|
||||
# The following line is equivalent to:
|
||||
# `output[vi, *vs] = storage[seq_id, history_id, *vs]`
|
||||
# However, unpacking operator in subscript requires Python 3.11 or newer
|
||||
T.buffer_store(
|
||||
output,
|
||||
T.BufferLoad(storage, [seq_id, history_id, *vs]),
|
||||
[vi, *vs],
|
||||
)
|
||||
|
||||
return f
|
||||
|
||||
return _func_one_dim() if len(shape) == 1 else _func_high_dim()
|
||||
|
||||
@staticmethod
|
||||
def create_set_func(
|
||||
shape: Sequence[Union[int, tirx.Var]],
|
||||
dtype: str,
|
||||
max_batch_size: Union[int, tirx.Var],
|
||||
max_history: Union[int, tirx.Var],
|
||||
state_id: int,
|
||||
) -> tirx.PrimFunc:
|
||||
"""Create the set function with given state shape.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape : Sequence[Union[int, tirx.Var]]
|
||||
The shape of the state tensor.
|
||||
|
||||
dtype: str
|
||||
The data type of the state tensor.
|
||||
|
||||
max_batch_size : Union[int, tirx.Var]
|
||||
The maximum batch size.
|
||||
|
||||
max_history : Union[int, tirx.Var]
|
||||
The maximum history length.
|
||||
|
||||
state_id : int
|
||||
The id of the state, used for naming the function.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tirx.PrimFunc
|
||||
The set function.
|
||||
"""
|
||||
|
||||
def _func_one_dim():
|
||||
@T.prim_func(s_tir=True)
|
||||
def f(
|
||||
var_storage: T.handle,
|
||||
var_seq_slot_ids: T.handle,
|
||||
var_history_slot_ids: T.handle,
|
||||
var_data: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
T.func_attr({"global_symbol": f"rnn_state_set_{state_id}"})
|
||||
|
||||
storage = T.match_buffer(
|
||||
var_storage, (max_batch_size, max_history, shape[0]), dtype
|
||||
)
|
||||
seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32")
|
||||
history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32")
|
||||
data = T.match_buffer(var_data, (batch_size, shape[0]), dtype)
|
||||
|
||||
for i in range(batch_size):
|
||||
for s in range(shape[0]):
|
||||
with T.sblock("copy"):
|
||||
vi, vs = T.axis.remap("SS", [i, s])
|
||||
seq_id: T.int32 = seq_slot_ids[vi]
|
||||
history_id: T.int32 = (history_slot_ids[vi] + 1) % T.cast(
|
||||
max_history, "int32"
|
||||
)
|
||||
storage[seq_id, history_id, vs] = data[vi, vs]
|
||||
|
||||
return f
|
||||
|
||||
def _func_high_dim():
|
||||
@T.prim_func(s_tir=True)
|
||||
def f(
|
||||
var_storage: T.handle,
|
||||
var_seq_slot_ids: T.handle,
|
||||
var_history_slot_ids: T.handle,
|
||||
var_data: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
T.func_attr({"global_symbol": f"rnn_state_set_{state_id}"})
|
||||
|
||||
storage = T.match_buffer(var_storage, (max_batch_size, max_history, *shape), dtype)
|
||||
seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32")
|
||||
history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32")
|
||||
data = T.match_buffer(var_data, (batch_size, *shape), dtype)
|
||||
|
||||
for i in range(batch_size):
|
||||
for s in T.grid(*shape):
|
||||
with T.sblock("copy"):
|
||||
vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s])
|
||||
seq_id: T.int32 = seq_slot_ids[vi]
|
||||
history_id: T.int32 = (history_slot_ids[vi] + 1) % T.cast(
|
||||
max_history, "int32"
|
||||
)
|
||||
# The following line is equivalent to:
|
||||
# `storage[seq_id, history_id, *vs] = data[vi, *vs]`
|
||||
# However, unpacking operator in subscript requires Python 3.11 or newer
|
||||
T.buffer_store(
|
||||
storage,
|
||||
T.BufferLoad(data, [vi, *vs]),
|
||||
[seq_id, history_id, *vs],
|
||||
)
|
||||
|
||||
return f
|
||||
|
||||
return _func_one_dim() if len(shape) == 1 else _func_high_dim()
|
||||
Reference in New Issue
Block a user