0x0. 前言
本文解析一下mlc-llm（https://github.com/mlc-ai/mlc-llm）对大模型推理的流程以及使用的图优化，算子优化策略。mlc-llm的模型部署流程可以查看官方文档：https://mlc.ai/mlc-llm/docs/ ，也可以参考我前段时间写的这篇MLC-LLM 部署RWKV World系列模型实战（3B模型Mac M2解码可达26tokens/s） 。
此外，阅读mlc-llm的代码还需要理解一些TVM Unify的一些基础概念，可以参考TVM 学习指南（个人版） ，Relax: TVM 的下一代图层级 IR，新一代深度学习编译技术变革和展望等等。从 https://github.com/BBuf/tvm_mlir_learn 这里可以查看更多相关博客和资料。
在 MLC-LLM 部署RWKV World系列模型实战（3B模型Mac M2解码可达26tokens/s） 中提到要使用mlc-llm部署模型首先需要一个编译过程，将原始的基于Realx搭建的模型比如RWKV和给定的device信息一起编译为TVM中的runtime.Module（在linux上编译的产物就是.so文件）提供mlc-llm的c++推理接口调用 。我们就从这里看起：
由于mlc-llm上游更新很快，为了准确标定代码位置我fork了一份2023年9月17号的mlc-llm代码 ：https://github.com/BBuf/mlc-llm-code-analysis ，本文的注释以及指出的代码位置均以这个fork仓库为准。
0x1. 编译流程解析
编译的入口在：https://github.com/BBuf/mlc-llm-code-analysis/blob/main/mlc_llm/build.py 。
这个脚本构建了一个模型build的入口，可以通过传入不同的参数来构建不同配置的模型。参数解析和模型编译都在 https://github.com/BBuf/mlc-llm-code-analysis/blob/main/mlc_llm/core.py 中实现，模型编译准备（mod_transform_before_build函数）和编译（build函数）两个阶段。在模型编译准备阶段，包含准备需要优化的算子，执行一些基础的图变换，针对cuda做进一步优化，做算子fuse等优化，详细的解释清阅读这里的注释：https://github.com/BBuf/mlc-llm-code-analysis/blob/main/mlc_llm/core.py#L378 。
在这之后会执行编译过程：https://github.com/BBuf/mlc-llm-code-analysis/blob/main/mlc_llm/core.py#L378 。从这里我们可以看到，对于GPU来说使用的是默认的schedule模板，并没有使用AutoTVM/Ansor等等调优工具，这一点是很友好的，个人猜测也是因为Transformer架构的模型是很固定的，然后优化方法也比较统一。
上面的编译前准备和编译都是针对IRModule来说的，那么这个IRModule是怎么来的呢？以及量化是在哪里做的？这两个问题都是在 build_model_from_args 函数： https://github.com/BBuf/mlc-llm-code-analysis/blob/main/mlc_llm/core.py#L627 处理的，发生在 mod_transform_before_build 函数调用之前。以 RWKV 模型为例，通过这行 mod, param_manager, params, model_config = rwkv.get_model(args, config) 代码完成了从原始的 huggingface 模型到初始的 IRModule 的转换，在这个过程中也包含了量化。
0x2. 模型搭建解析
0x2.1 模型组件搭建
首先在 https://github.com/BBuf/mlc-llm-code-analysis/blob/main/mlc_llm/relax_model/modules.py 这里基于Relax的内部接口（relax.Expr，relax.testing.nn.Module，relax.op.xxx等等）定义了搭建LLM模型需要的一些组件比如 ModuleList，Linear，Embedding，LayerNorm，RotaryEmbedding等等。这个地方我添加了一些解释，请点上面的源码链接查看。然后这个地方需要注意2个特殊的op，第一个是来自 https://github.com/mlc-ai/relax/blob/ceaf7b0156524d30537a3de5fa30764eaff4edb8/python/tvm/relax/op/index.py#L28 的：
def take(x: Expr, indices: Expr, axis: Optional[int] = None) -> Expr:
return _ffi_api.take(x, indices, axis) # type: ignore
这个函数，实现了take的核心功能，与numpy和pytorch的take语义类似,都可以通过指定indices来从输入张量中抽取值。主要调用了_ffi_api.take进行取值操作, 这个_ffi_api是relax底层实现, take操作的实际计算会在这里进行。这个函数被用于Embedding组件的搭建中。
另外nn.emit这个接口的作用是将一个relax.Expr表达式转化为relax.Var变量，并保存该变量。
最后我们注意到这里搭建的Relax模块风格和PyTorch的模块风格基本一致，也可以看出Relax前端是不断靠近动态图风格，追求更佳的易用性。
0x2.2 模型搭建
首先看一些准备工作：
# @dataclass：这个装饰器用于指示RWKVConfig类是一个数据类。用于存储RWKVModel的配置信息。
@dataclass
class RWKVConfig:
"""The configuration class to store the configuration of a `RWKVModel`."""
num_hidden_layers: int # 类中的一个属性，用于存储隐藏层的数量，类型为整数。
vocab_size: int # 类中的一个属性，用于存储词汇表的大小，类型为整数。
hidden_size: int # 类中的一个属性，用于存储隐藏层的大小，类型为整数。
intermediate_size: int # 类中的一个属性，用于存储中间层的大小，类型为整数。
rescale_every: int = 0 # 类中的一个属性，默认值为0，用于存储重新缩放的频率，类型为整数。
layer_norm_epsilon: float = 1e-5 # 类中的一个属性，默认值为1e-5，用于存储层归一化的epsilon值，类型为浮点数。
max_sequence_length: int = 1024 # 类中的一个属性，默认值为1024，用于存储最大序列长度，类型为整数。
dtype: str = "float32" # 类中的一个属性，默认值为"float32"，用于存储数据类型，类型为字符串。
def __init__(
self,
num_hidden_layers: int,
vocab_size: int,
hidden_size: int,
intermediate_size: int,
rescale_every: int = 0,
layer_norm_epsilon: float = 1e-5,
context_length: int = 1024,
dtype: str = "float32",
**kwargs,
) -> None:
self.num_hidden_layers = num_hidden_layers
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.rescale_every = rescale_every
self.layer_norm_epsilon = layer_norm_epsilon
self.max_sequence_length = context_length
self.dtype = dtype
self.kwargs = kwargs
# 用来索引RWKV的Attention和FFN部分存储的状态或者叫Cache。
# python代码可以参考： https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L858-L867
class State:
ATT_X = 0
ATT_A = 1
ATT_B = 2
ATT_P = 3
FFN_X = 4
这里的State是用来索引RWKV的Attention和FFN部分存储的状态或者叫Cache，每一个Layer有5个不同的State，并且每个State的shape都是[1, hidden_size]，这里的1代表的应该是batch纬度。
# 义了一个名为_load_state的函数，它接受一个名为state的参数，类型为Expr，一个名为hidden_size的参数，类型为整数，
# 一个名为dtype的参数，类型为字符串。函数的返回类型为Expr。
def _load_state(state: Expr, hidden_size: int, dtype: str) -> Expr:
# Reuse `attention_kv_cache_view`
# 将外部函数vm.builtin.attention_kv_cache_view赋值给变量f_load_cache。relax.extern是一个外部函数调用的语法，
# 它指示编译器在编译时将该函数调用转换为相应的外部函数调用。
f_load_cache = relax.extern("vm.builtin.attention_kv_cache_view")
# 使用nn.emit方法生成一个表达式对象，该表达式表示对外部函数f_load_cache的调用。
# 调用的参数是一个列表，包含state和R.shape([1, hidden_size])，以及sinfo_args参数指定的一个R.Tensor对象。
cache = nn.emit(
relax.Call(
f_load_cache,
[state, R.shape([1, hidden_size])],
sinfo_args=[R.Tensor((1, hidden_size), dtype)],
)
)
return cache
# 定义了一个名为_store_state的函数，它接受一个名为state的参数，类型为Expr，一个名为value的参数，类型为Expr。
def _store_state(state: Expr, value: Expr):
# Reuse `attention_kv_cache_update`
# 将外部函数vm.builtin.attention_kv_cache_update赋值给变量f_store_cache。
# relax.extern是一个外部函数调用的语法，它指示编译器在编译时将该函数调用转换为相应的外部函数调用。
f_store_cache = relax.extern("vm.builtin.attention_kv_cache_update")
# 使用nn.emit方法生成一个表达式对象，该表达式表示对外部函数f_store_cache的调用。
# 调用的参数是一个列表，包含state和value，以及sinfo_args参数指定的一个R.Object()对象。
return nn.emit(
relax.Call(
f_store_cache,
[state, value],
sinfo_args=[R.Object()],
)
)
这两个函数用来加载和存储RWKV模型的State。接下来看一下对应 https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L741 这里的torch.ops.rwkv.wkv_forward(1, T, C, w, u, k, v, y, aa, bb, pp) 的Relax实现，为了方便对照先贴一下原始的wkv forward cuda kernel：
template <typename F>
__global__ void kernel_wkv_forward(const int B, const int T, const int C,
const float *__restrict__ const _w, const float *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v,
F *__restrict__ const _y, float *__restrict__ const _aa, float *__restrict__ const _bb, float *__restrict__ const _pp) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int _b = idx / C;
const int _c = idx % C;
const int _offset = _b * T * C + _c;
const int _state_offset = _b * C + _c;
float u = _u[_c];
float w = _w[_c];
const F *__restrict__ const k = _k + _offset;
const F *__restrict__ const v = _v + _offset;
F *__restrict__ const y = _y + _offset;
float aa = _aa[_state_offset];
float bb = _bb[_state_offset];
float pp = _pp[_state_offset];
for (int i = 0; i < T; i++) {
const int ii = i * C;
const float kk = float(k[ii]);
const float vv = float(v[ii]);
float ww = u + kk;
float p = max(pp, ww);
float e1 = exp(pp - p);
float e2 = exp(ww - p);
y[ii] = F((e1 * aa + e2 * vv) / (e1 * bb + e2));
ww = w + pp;
p = max(ww, kk);
e1 = exp(ww - p);
e2 = exp(kk - p);
aa = e1 * aa + e2 * vv;
bb = e1 * bb + e2;
pp = p;
}
_aa[_state_offset] = aa;
_bb[_state_offset] = bb;
_pp[_state_offset] = pp;
}
template <typename F>
void cuda_wkv_forward(int B, int T, int C, float *w, float *u, F *k, F *v, F *y, float *aa, float *bb, float *pp) {
dim3 threadsPerBlock( min(C, 32) );
assert(B * C % threadsPerBlock.x == 0);
dim3 numBlocks(B * C / threadsPerBlock.x);
kernel_wkv_forward<<<numBlocks, threadsPerBlock>>>(B, T, C, w, u, k, v, y, aa, bb, pp);
}
这个cuda kernel里面，B表示batch_size，在mlc-llm的实现默认为1。然后T表示序列长度，C表示隐藏层纬度。然后我们就可以对应来看mlc-llm的wkv实现了。
# 定义了一个名为create_wkv_func的函数，它接受一个名为hidden_size的参数，
# 类型为整数，一个名为dtype的参数，类型为字符串，一个名为out_dtype的参数，类型为字符串。
def create_wkv_func(hidden_size: int, dtype: str, out_dtype: str):
@T.prim_func
def wkv_func(
k: T.handle,
v: T.handle,
time_decay: T.handle,
time_first: T.handle,
saved_a: T.handle,
saved_b: T.handle,
saved_p: T.handle,
wkv: T.handle,
out_a: T.handle,
out_b: T.handle,
out_p: T.handle,
):
# 设置TIR函数的属性。这里设置了三个属性，包括op_pattern、tir.noalias和tir.is_scheduled。
T.func_attr({"op_pattern": 8, "tir.noalias": True, "tir.is_scheduled": 1})
# 声明一个名为context_length的变量，类型为T.int64()，用于存储上下文长度。
context_length = T.int64()
# 创建一个名为K的匹配缓冲区，通过T.match_buffer方法匹配参数k的形状和数据类型。
# K的形状在原始的ChatRWKV中为B，T，C，只不过这里B=1
# 这里的k就是上面cuda kernel的_k
K = T.match_buffer(k, (context_length, hidden_size), dtype=dtype)
# 创建一个名为V的匹配缓冲区，通过T.match_buffer方法匹配参数v的形状和数据类型。
# 这里的v就是上面cuda kernel的_v
V = T.match_buffer(v, (context_length, hidden_size), dtype=dtype)
# 创建一个名为TimeDecay的匹配缓冲区，通过T.match_buffer方法匹配参数time_decay的形状和数据类型。
# 这里的TimeDecay就是上面的w
TimeDecay = T.match_buffer(time_decay, (hidden_size,), dtype=dtype)
# 创建一个名为TimeFirst的匹配缓冲区，通过T.match_buffer方法匹配参数time_first的形状和数据类型。
# 这里的TimeFirst对应上面的u
TimeFirst = T.match_buffer(time_first, (hidden_size,), dtype=dtype)
# 对应kernel里面的_aa的上一个token的状态
SavedA = T.match_buffer(saved_a, (1, hidden_size), dtype=dtype)
# 对应kernel里面的_bb的上一个token的状态
SavedB = T.match_buffer(saved_b, (1, hidden_size), dtype=dtype)
# 对应kernel里面的_pp的上一个token的状态
SavedP = T.match_buffer(saved_p, (1, hidden_size), dtype=dtype)
# 对应_aa的当前token状态
OutA = T.match_buffer(out_a, (1, hidden_size), dtype=dtype)
# 对应_bb的当前token状态
OutB = T.match_buffer(out_b, (1, hidden_size), dtype=dtype)
# 对应_pp的当前token状态
OutP = T.match_buffer(out_p, (1, hidden_size), dtype=dtype)
# 对应kernel里面的p
P = T.alloc_buffer((hidden_size,), dtype=dtype, scope="local")
# 对应kernel里面的e1
E1 = T.alloc_buffer((hidden_size,), dtype=dtype, scope="local")
# 对应kernel里面的e2
E2 = T.alloc_buffer((hidden_size,), dtype=dtype, scope="local")
# 对应kernel里面的aa
A_local = T.alloc_buffer((hidden_size,), dtype=dtype, scope="local")
# 对应kernel里面的bb
B_local = T.alloc_buffer((hidden_size,), dtype=dtype, scope="local")
# 对应kernel里面的cc
P_local = T.alloc_buffer((hidden_size,), dtype=dtype, scope="local")
# 迭代hidden_size // 32次，使用T.thread_binding方法进行线程绑定，其中hidden_size // 32是块索引的范围。
# 这里的线程块划分和rwkv kernel里面保持一致：即每个block 32个线程，一共((B=1)*C)/32个blcok
for bx in T.thread_binding(hidden_size // 32, thread="blockIdx.x"):
# 迭代32次，使用T.thread_binding方法进行线程绑定，其中32是线程索引的范围。
for tx in T.thread_binding(32, thread="threadIdx.x"):
# 创建一个名为"init"的块，用于初始化局部变量。
with T.block("init"):
# 对应 const int _state_offset = _b * C + _c;
vi = T.axis.S(hidden_size, bx * 32 + tx)
# 对应 float aa = _aa[_state_offset];
A_local[vi] = SavedA[0, vi]
# 对应 float bb = _bb[_state_offset];
B_local[vi] = SavedB[0, vi]
# 对应 float pp = _pp[_state_offset];
P_local[vi] = SavedP[0, vi]
for j in range(context_length): # 对应 for (int i = 0; i < T; i++)
with T.block("main"):
# 对应 const int _state_offset = _b * C + _c;
vi = T.axis.S(hidden_size, bx * 32 + tx)
# vj 对应 _b * T; [vj, vi] = _b * T * C + _b * C + _c
# _b * T * C + _c = _offset
vj = T.axis.opaque(context_length, j)
# 对应 float p = max(pp, ww); float ww = u + kk;
# const float kk = float(k[ii]); const int ii = i * C;
# const F *__restrict__ const k = _k + _offset;
P[vi] = T.max(P_local[vi], K[vj, vi] + TimeFirst[vi])
# 对应 float e1 = exp(pp - p);
E1[vi] = T.exp(P_local[vi] - P[vi])
# 对应 float e2 = exp(ww - p);
E2[vi] = T.exp(K[vj, vi] + TimeFirst[vi] - P[vi])
P[vi] = T.max(P_local[vi] + TimeDecay[vi], K[vj, vi])
E1[vi] = T.exp(P_local[vi] + TimeDecay[vi] - P[vi])
E2[vi] = T.exp(K[vj, vi] - P[vi])
A_local[vi] = E1[vi] * A_local[vi] + E2[vi] * V[vj, vi]
B_local[vi] = E1[vi] * B_local[vi] + E2[vi]
P_local[vi] = P[vi]
with T.block("write_back"):
vi = T.axis.S(hidden_size, bx * 32 + tx) # 对应
OutA[0, vi] = A_local[vi] # 对应 _aa[_state_offset] = aa;
OutB[0, vi] = B_local[vi] # 对应 _bb[_state_offset] = bb;
OutP[0, vi] = P_local[vi] # 对应 _pp[_state_offset] = pp;
return wkv_func
我们可以看到mlc-llm里面的wkv forward实现基本就是用基于Relax的api将cuda函数翻译成了TIR。注释里面给了一些下标的推导以及每一行Relax的代码是如何对应到原始的cuda kernel。
# 定义了一个名为_te_concat_saved_x的函数，它接受两个参数saved_x和x，都是te.Tensor类型的张量。
# 使用TVM的te.compute函数计算一个新的张量，该张量的形状与x相同，元素根据条件判断进行选择。如果i等于0，
# 则选择saved_x[0, j]作为元素值，否则选择x[i - 1, j]作为元素值。其中i和j是迭代变量。
def _te_concat_saved_x(saved_x: te.Tensor, x: te.Tensor):
return te.compute(
x.shape,
lambda i, j: tir.if_then_else(i == 0, saved_x[0, j], x[i - 1, j]),
)
# 定义了一个名为_te_get_last_x的函数，它接受一个参数x，是一个te.Tensor类型的张量。
# a. seq_len, hidden_size = x.shape：获取x张量的形状，其中seq_len表示序列长度，hidden_size表示隐藏大小。
# b. return te.compute(...)：使用TVM的te.compute函数计算一个新的张量，该张量的形状为(1, hidden_size)，
# 元素值为x[seq_len - 1, j]，其中j是迭代变量。
def _te_get_last_x(x: te.Tensor):
seq_len, hidden_size = x.shape
return te.compute((1, hidden_size), lambda _, j: x[seq_len - 1, j])
这两个函数应该对应了 https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L455 这里代码里面的sx = torch.cat((sx.unsqueeze(0), xx[:-1,:]))和xx[-1, :]：
@MyFunction
def ffn_seq(self, x, sx, ln_w, ln_b, k_mix, r_mix, kw, vw, rw, kmx, krx, kmy, kry, vmx, vrx, vmy, vry, rmx, rrx, rmy, rry):
xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b)
sx = torch.cat((sx.unsqueeze(0), xx[:-1,:]))
kx = xx * k_mix + sx * (1 - k_mix)
rx = xx * r_mix + sx * (1 - r_mix)
r = torch.sigmoid(gemm(rx, rw))
vx = torch.square(torch.relu(gemm(kx, kw)))
out = r * gemm(vx, vw)
return x + out, xx[-1,:]
接着对Embedding函数进行解析：
# 定义了一个名为RWKV_Embedding的PyTorch模块。
class RWKV_Embedding(nn.Module):
# 定义了RWKV_Embedding类的构造函数，接受三个参数num_embeddings、embedding_dim和dtype。
def __init__(self, num_embeddings, embedding_dim, dtype):
self.num_embeddings = num_embeddings # 将num_embeddings赋值给类成员变量self.num_embeddings。
self.embedding_dim = embedding_dim # 将embedding_dim赋值给类成员变量self.embedding_dim。
# 创建一个名为weight的Parameter，形状为(num_embeddings, embedding_dim)，
# 数据类型为dtype，并将其赋值给类成员变量self.weight。
self.weight = nn.Parameter(
(num_embeddings, embedding_dim), dtype=dtype, name="weight"
)
def forward(self, x: relax.Expr) -> relax.Var:
# 调用op.reshape函数将输入张量x进行reshape，将其展平为一维张量，并将结果重新赋值给x。
# nn.emit是将一个relax.Expr表达式转化为relax.Var变量，并保存该变量。
x = nn.emit(op.reshape(x, shape=[-1]))
# 使用op.take操作从self.weight中按照索引x提取对应的嵌入向量，并返回结果。这里的axis=0表示在第一个维度上进行索引操作。
return nn.emit(op.take(self.weight, x, axis=0))
以及LayerNorm：
# 这段代码定义了一个名为RWKV_LayerNorm的PyTorch模块，它实现了一个Layer Normalization层。
class RWKV_LayerNorm(nn.Module):
# 定义了RWKV_LayerNorm类的构造函数，接受四个参数intermediate_size、dtype、eps和name_prefix。
def __init__(self, intermediate_size, dtype, eps=1e-5, name_prefix=""):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(
(intermediate_size,), dtype=dtype, name=f"{name_prefix}_ln_weight"
)
self.bias = nn.Parameter(
(intermediate_size,), dtype=dtype, name=f"{name_prefix}_ln_bias"
)
def forward(self, x: relax.Expr) -> relax.Var:
# 使用op.nn.layer_norm操作对输入张量x进行Layer Normalization，其中使用Parameter self.weight作为缩放参数（gamma），
# 使用可学习参数self.bias作为偏移参数（beta），在最后一个维度（axes=-1）上进行标准化操作，
# 并设置小数值修正项为self.eps。将标准化后的结果重新赋值给x。
x = nn.emit(
op.nn.layer_norm(
x,
gamma=self.weight,
beta=self.bias,
axes=-1,
epsilon=self.eps,
)
)
return x
接着对FFN层做一个详细的解析：
# 这段代码定义了一个名为RWKV_FFN的PyTorch模块，它实现了Feed-Forward Network（FFN）。
class RWKV_FFN(nn.Module):
# 定义了RWKV_FFN类的构造函数，接受两个参数RWKVConfig和index。
def __init__(self, config: RWKVConfig, index: int) -> None:
super().__init__()
# 将config.hidden_size赋值给类成员变量self.hidden_size，表示隐藏大小。
self.hidden_size = config.hidden_size
# 将config.dtype赋值给类成员变量self.dtype，表示数据类型。
self.dtype = config.dtype
# 将index赋值给类成员变
self.index = index
# 建一个名为time_mix_key的可学习参数，形状为(self.hidden_size,)，
# 数据类型为config.dtype，命名为"ffn_{index}_time_mix_k"，并将其赋值给类成员变量self.time_mix_key。
self.time_mix_key = nn.Parameter(
(self.hidden_size,), dtype=config.dtype, name=f"ffn_{index}_time_mix_k"
)
# 创建一个名为time_mix_receptance的可学习参数，形状为(self.hidden_size,)，数据类型为config.dtype，
# 命名为"ffn_{index}_time_mix_r"，并将其赋值给类成员变量self.time_mix_receptance。
self.time_mix_receptance = nn.Parameter(
(self.hidden_size,), dtype=config.dtype, name=f"ffn_{index}_time_mix_r"
)
# 创建一个线性层，输入大小为self.hidden_size，输出大小为config.intermediate_size，
# 数据类型为config.dtype，没有偏置项，并将其赋值给类成员变量self.key。
self.key = Linear(
self.hidden_size, config.intermediate_size, dtype=config.dtype, bias=False
)
# 创建一个线性层，输入大小为self.hidden_size，输出大小为self.hidden_size，数据类型为config.dtype，
# 没有偏置项，并将其赋值给类成员变量self.receptance。
self.receptance = Linear(
self.hidden_size, self.hidden_size, dtype=config.dtype, bias=False
)
self.value = Linear(
config.intermediate_size, self.hidden_size, dtype=config.dtype, bias=False
)
def forward(self, x: Expr, state: Expr) -> Expr:
# 计算偏移量，用于在state中获取对应的保存状态。
offset = self.index * 5 + State.FFN_X
# 获取x的shape[0]表示上下文长度。
context_length = x.struct_info.shape[0]
# 获取隐藏层大小。
hidden_size = self.hidden_size
# 调用_load_state函数从state中加载保存的状态state[offset]，并将结果赋值给saved_x。
saved_x = _load_state(state[offset], hidden_size, self.dtype)
# 如果上下文长度不为1，则执行下面的操作。
if not is_one(context_length):
# 调用nn.emit_te函数，将saved_x和x作为参数传递给
# _te_concat_saved_x函数进行计算，并将结果重新赋值给saved_x。
# 类似于transformer 里面的KV Cache的，但是这里的concat是纬度不变的
# 对应 sx = torch.cat((sx.unsqueeze(0), xx[:-1,:])) 这行代码
saved_x = nn.emit_te(_te_concat_saved_x, saved_x, x)
# 创建一个全为1的张量，形状为(hidden_size,)，数据类型为self.dtype，并将其赋值给ones。
ones = nn.emit(relax.op.ones((hidden_size,), self.dtype))
# 计算xk，根据时间混合参数self.time_mix_key和保存的状态saved_x，使用加权求和的方式得到。
# 其中，x和saved_x分别乘以self.time_mix_key和(ones - self.time_mix_key)，然后相加。将计算结果赋值给xk。
# 对应 kx = xx * k_mix + sx * (1 - k_mix) 这行代码
xk = nn.emit(x * self.time_mix_key + saved_x * (ones - self.time_mix_key))
# 计算xr，根据时间混合参数self.time_mix_receptance和保存的状态saved_x，使用加权求和的方式得到。
# 其中，x和saved_x分别乘以self.time_mix_receptance和(ones - self.time_mix_receptance)，然后相加。
# 将计算结果赋值给xr。
# 对应 rx = xx * r_mix + sx * (1 - r_mix)
xr = nn.emit(
x * self.time_mix_receptance + saved_x * (ones - self.time_mix_receptance)
)
# # 如果上下文长度不为1，则执行下面的操作。
if not is_one(context_length):
# 调用nn.emit_te函数，使用_te_get_last_x函数从x中获取最后一个token对应的tensor，并将结果重新赋值给x。
# 对应 xx[-1,:]
x = nn.emit_te(_te_get_last_x, x)
# 断言x的结构信息（shape）的第一个维度为1。
assert is_one(x.struct_info.shape[0])
# 调用_store_state函数，将x保存到state[offset]中，并将结果重新赋值给saved_x。
# 对应：https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L921
saved_x = _store_state(state[offset], x)
# 将xr作为输入，经过sigmoid激活函数计算得到r。对应：r = torch.sigmoid(gemm(rx, rw))
r = nn.emit(op.sigmoid(self.receptance(xr)))
# 对应 vx = torch.square(torch.relu(gemm(kx, kw)))
xv = nn.emit(op.square(op.nn.relu(self.key(xk))))
return nn.emit(r * self.value(xv)), [saved_x]
接下来对Attention部分的实现进行解析，注意这部分对应的代码在 https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L728-L747 。贴一下python代码防止看错位置产生疑问：
if os.environ["RWKV_CUDA_ON"] == '1':
@MyFunction
def cuda_att_seq(self, x, sx, aa, bb, pp, ln_w, ln_b, k_mix, v_mix, r_mix, t_decay, t_first, kw, vw, rw, ow, kmx, krx, kmy, kry, vmx, vrx, vmy, vry, rmx, rrx, rmy, rry, omx, orx, omy, ory):
T, C = x.shape
xx = F.layer_norm(x, (C,), weight=ln_w, bias=ln_b)
sx = torch.cat((sx.unsqueeze(0), xx[:-1,:]))
kx = xx * k_mix + sx * (1 - k_mix)
vx = xx * v_mix + sx * (1 - v_mix)
rx = xx * r_mix + sx * (1 - r_mix)
r = torch.sigmoid(gemm(rx, rw))
k = gemm(kx, kw, output_dtype=torch.float32)
v = gemm(vx, vw, output_dtype=torch.float32)
y, aa, bb, pp = cuda_wkv(T, aa.shape[0], t_decay, t_first, k, v, aa, bb, pp)
out = gemm(r * y.to(x.dtype), ow)
return x + out, xx[-1,:], aa, bb, pp
对应mlc-llm RWKV Attention的代码解析为：
# 实现RWKV Attention，对应 https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L479
class RWKV_Attention(nn.Module):
# 初始化函数，接受一个config对象和一个整数index作为参数。其中config是一个RWKVConfig类型的对象，index表示当前层的索引。
def __init__(self, config: RWKVConfig, index: int) -> None:
super().__init__()
self.index = index
self.dtype = config.dtype
self.hidden_size = config.hidden_size
# 创建一些可学习的参数，如time_decay、time_first、time_mix_key等，这些参数会在模型的前向传播中使用。
self.time_decay = nn.Parameter(
(self.hidden_size,), dtype="float32", name=f"att_{index}_time_decay"
)
self.time_first = nn.Parameter(
(self.hidden_size,), dtype="float32", name=f"att_{index}_time_first"
)
self.time_mix_key = nn.Parameter(
(self.hidden_size,), dtype=config.dtype, name=f"att_{index}_time_mix_k"
)
self.time_mix_value = nn.Parameter(
(self.hidden_size,), dtype=config.dtype, name=f"att_{index}_time_mix_v"
)
self.time_mix_receptance = nn.Parameter(
(self.hidden_size,), dtype=config.dtype, name=f"att_{index}_time_mix_r"
)
# 前向传播用到的线性层
self.key = Linear(
self.hidden_size, self.hidden_size, dtype=config.dtype, bias=False
)
self.value = Linear(
self.hidden_size, self.hidden_size, dtype=config.dtype, bias=False
)
self.receptance = Linear(
self.hidden_size, self.hidden_size, dtype=config.dtype, bias=False
)
self.output = Linear(
self.hidden_size, self.hidden_size, dtype=config.dtype, bias=False
)
# 前向传播函数，接受输入张量x和状态张量state作为参数，并返回输出张量
def forward(self, x: Expr, state: Expr) -> Expr:
# Load current state
# 定义了一些局部变量，如ones、index、hidden_size、context_length等。
ones = nn.emit(relax.op.ones((self.hidden_size,), self.dtype))
index = self.index
hidden_size = self.hidden_size
context_length = x.struct_info.shape[0]
bb = relax.BlockBuilder.current()
# _load_state函数从state中加载保存的状态，赋值给saved_a、saved_b、saved_p和saved_x。
saved_a = _load_state(state[index * 5 + State.ATT_A], hidden_size, "float32")
saved_b = _load_state(state[index * 5 + State.ATT_B], hidden_size, "float32")
saved_p = _load_state(state[index * 5 + State.ATT_P], hidden_size, "float32")
saved_x = _load_state(state[index * 5 + State.ATT_X], hidden_size, self.dtype)
# 调用nn.emit_te函数，将saved_x和x作为参数传递给
# _te_concat_saved_x函数进行计算，并将结果重新赋值给saved_x。
# 对应 sx = torch.cat((sx.unsqueeze(0), xx[:-1,:]))
if not is_one(context_length):
saved_x = nn.emit_te(_te_concat_saved_x, saved_x, x)
# 对应 kx = xx * k_mix + sx * (1 - k_mix)
xk = nn.emit(x * self.time_mix_key + saved_x * (ones - self.time_mix_key))
# 对应 vx = xx * v_mix + sx * (1 - v_mix)
xv = nn.emit(x * self.time_mix_value + saved_x * (ones - self.time_mix_value))
# 对应 rx = xx * r_mix + sx * (1 - r_mix)
xr = nn.emit(
x * self.time_mix_receptance + saved_x * (ones - self.time_mix_receptance)
)
# 对应 r = torch.sigmoid(gemm(rx, rw))
r = nn.emit(op.sigmoid(self.receptance(xr)))
# 对应 k = gemm(kx, kw, output_dtype=torch.float32)
k = nn.emit(op.astype(self.key(xk), "float32"))
# 对应 v = gemm(vx, vw, output_dtype=torch.float32)
v = nn.emit(op.astype(self.value(xv), "float32"))
# 这部分对应 y, aa, bb, pp = cuda_wkv(T, aa.shape[0], t_decay, t_first, k, v, aa, bb, pp)
# 这里的 create_wkv_func 在上面已经解析了
gv = bb.add_func(create_wkv_func(hidden_size, "float32", self.dtype), "wkv")
ret = nn.emit(
relax.call_tir(
gv,
[k, v, self.time_decay, self.time_first, saved_a, saved_b, saved_p],
[
R.Tensor((context_length, hidden_size), self.dtype), # 对应wkv
R.Tensor((1, hidden_size), "float32"), # 对应out_a
R.Tensor((1, hidden_size), "float32"), # 对应out_b
R.Tensor((1, hidden_size), "float32"), # 对应out_p
],
)
)
if not is_one(context_length):
# 对应 xx[-1,:]
x = nn.emit_te(_te_get_last_x, x)
assert is_one(x.struct_info.shape[0])
saved_x = _store_state(state[self.index * 5 + State.ATT_X], x)
saved_a = _store_state(state[self.index * 5 + State.ATT_A], ret[1])
saved_b = _store_state(state[self.index * 5 + State.ATT_B], ret[2])
saved_p = _store_state(state[self.index * 5 + State.ATT_P], ret[3])
# 需要注意一下，python代码里面的 return x + out, xx[-1,:], aa, bb, pp
# 这里的 x + out被放在attention外面做了，因为这里的x已经是被修改之后好的结果而不是原始的x
return nn.emit(self.output(r * ret[0])), [
saved_x,
saved_a,
saved_b,
saved_p,
]
接着解析一下RWKVLayer的实现，请注意下面的最后一行代码的解释：
class RWKVLayer(nn.Module):
# 初始化函数，接受一个config对象和一个整数index作为参数。其中config是一个RWKVConfig类型的对象，index表示层的索引。
def __init__(self, config: RWKVConfig, index: int) -> None:
super().__init__()
# 如果index为0，创建一个RWKV_LayerNorm对象pre_ln，用于对输入进行Layer Normalization操作。
if index == 0:
self.pre_ln = RWKV_LayerNorm(
config.hidden_size,
config.dtype,
eps=config.layer_norm_epsilon,
name_prefix="pre_ln",
)
# 创建两个RWKV_LayerNorm对象，分别命名为ln1和ln2，
# 用于对注意力机制和前馈神经网络的输出进行Layer Normalization操作。
self.ln1 = RWKV_LayerNorm(
config.hidden_size,
config.dtype,
eps=config.layer_norm_epsilon,
name_prefix=f"att_{index}",
)
self.ln2 = RWKV_LayerNorm(
config.hidden_size,
config.dtype,
eps=config.layer_norm_epsilon,
name_prefix=f"ffn_{index}",
)
# 创建一个RWKV_Attention对象attention，用于实现注意力机制。
self.attention = RWKV_Attention(config, index)
# 创建一个RWKV_FFN对象feed_forward，用于实现前馈神经网络。
self.feed_forward = RWKV_FFN(config, index)
self.rescale_every = config.rescale_every
self.dtype = config.dtype
self.index = index
# 前向传播函数，接受输入张量x和状态张量state作为参数，并返回输出张量和更新后的状态列表。
def forward(self, x: Expr, state: Expr) -> Tuple[Expr, List[Expr]]:
# 如果index为0，则将输入张量x传入pre_ln进行Layer Normalization操作。
if self.index == 0:
x = self.pre_ln(x)
# 将经过ln1的输入张量x和状态张量state传入attention进行计算，得到注意力机制的输出att和更新后的状态列表att_state。
att, att_state = self.attention(self.ln1(x), state)
# 将输入张量x和注意力机制的输出att相加，并将结果赋值给x。
x = nn.emit(x + att)
# 将经过ln2的输入张量x和状态张量state传入feed_forward进行计算，得到前馈神经网络的输出ffn和更新后的状态列表ffn_state。
ffn, ffn_state = self.feed_forward(self.ln2(x), state)
# 将输入张量x和前馈神经网络的输出ffn相加，并将结果赋值给x。
x = nn.emit(x + ffn)
# 如果满足self.rescale_every > 0且(self.index + 1) % self.rescale_every == 0，则对输入张量x进行缩放操作。
if self.rescale_every > 0 and (self.index + 1) % self.rescale_every == 0:
x = nn.emit(x / relax.const(2, dtype=self.dtype))
# 返回输出张量x和注意力机制和前馈神经网络的更新后的状态列表的拼接。
return x, att_state + ffn_state
注意这里的attn_state是[saved_x,
saved_a,
saved_b,
saved_p,] ，然后ffn_state是[saved_x]，注意这两个x是不一样的，这5个状态也和本节开头的class State的成员定义一致。
接下来对RWKV模型定义进行了解析：
# 该代码是一个自定义的PyTorch模型类RWKVModel，继承自nn.Module
class RWKVModel(nn.Module):
# 初始化函数，接受一个config对象作为参数。其中config是一个RWKVConfig类型的对象。
def __init__(self, config: RWKVConfig) -> None:
super().__init__()
# 创建一个RWKV_Embedding对象embeddings，用于实现输入的嵌入操作。
self.embeddings = RWKV_Embedding(
num_embeddings=config.vocab_size,
embedding_dim=config.hidden_size,
dtype=config.dtype,
)
# 创建一个ModuleList对象blocks，其中包含了config.num_hidden_layers个RWKVLayer对象，
# 每个对象的索引从0到config.num_hidden_layers-1。
self.blocks = ModuleList(
[RWKVLayer(config, i) for i in range(config.num_hidden_layers)]
)
# 创建一个RWKV_LayerNorm对象ln_out，用于对输出进行Layer Normalization操作。
self.ln_out = RWKV_LayerNorm(
config.hidden_size,
config.dtype,
eps=config.layer_norm_epsilon,
name_prefix="out_ln",
)
self.hidden_size = config.hidden_size
self.dtype = config.dtype
# 前向传播函数，接受输入张量input_ids和状态张量state作为参数，并返回输出张量和更新后的状态列表。
def forward(self, input_ids: Expr, state: Expr) -> Tuple[Expr, List[Expr]]:
# 将输入张量input_ids传入embeddings进行嵌入操作，得到隐藏状态张量hidden_states。
hidden_states = self.embeddings(input_ids)
# 创建一个空列表states，用于存储每个RWKVLayer对象的更新后的状态列表。
states = []
# 遍历blocks中的每个RWKVLayer对象，将隐藏状态张量hidden_states和状态张量state传入
# 每个RWKVLayer对象的前向传播函数进行计算，得到更新后的隐藏状态张量和更新后的状态列表，
# 并将更新后的状态列表添加到states中。
for _, layer in enumerate(self.blocks):
hidden_states, layer_states = layer(hidden_states, state)
states += layer_states
# 获取隐藏状态张量的上下文长度context_length。
context_length = hidden_states.struct_info.shape[0]
# 如果context_length不为1，则调用_te_get_last_x函数获取最后一个token对应的张量。
if not is_one(context_length):
hidden_states = nn.emit_te(_te_get_last_x, hidden_states)
# 将隐藏状态张量传入ln_out进行Layer Normalization操作。
hidden_states = self.ln_out(hidden_states)
# 返回输出隐藏状态张量和所有RWKVLayer对象的更新后的状态列表。
return hidden_states, states
# 该代码是一个自定义的PyTorch模型类RWKVForCausalLM，继承自nn.Module。
class RWKVForCausalLM(nn.Module):
# 初始化函数，接受一个config对象作为参数。其中config是一个RWKVConfig类型的对象。
def __init__(self, config: RWKVConfig):
# 创建一个RWKVModel对象rwkv，用于实现序列模型的计算。
self.rwkv = RWKVModel(config)
# 创建一个Linear对象head，用于将隐藏状态映射到词汇表大小的输出空间。
self.head = Linear(
config.hidden_size, config.vocab_size, dtype=config.dtype, bias=False
)
self.vocab_size = config.vocab_size
############ End ############
# 前向传播函数，接受输入张量input_ids和状态张量state作为参数，并返回预测的logits和更新后的kv cache。
def forward(
self,
input_ids: relax.Expr,
state: relax.Expr,
):
# 将输入张量input_ids和状态张量state传入rwkv对象的前向传播函数进行计算，
# 得到更新后的隐藏状态张量hidden_states和key-value缓存key_value_cache。
hidden_states, key_value_cache = self.rwkv(input_ids, state)
# 将隐藏状态张量hidden_states传入head进行线性映射操作，得到logits。
logits = nn.emit(self.head(hidden_states))
# 对logits进行形状重塑，将其reshape为形状为(1, 1, self.vocab_size)的张量。
logits = nn.emit(op.reshape(logits, (1, 1, self.vocab_size)))
# 如果logits的数据类型不是float32，则将其转换为float32类型。
if logits.struct_info.dtype != "float32":
logits = nn.emit(relax.op.astype(logits, "float32"))
return logits, key_value_cache
解下是一个根据参数的名字确定量化参数类型的函数：
# 该代码定义了一个函数get_param_quant_kind，用于根据参数名称和参数信息确定参数的量化类型。
def get_param_quant_kind(
name: str, param_info: relax.TensorStructInfo
) -> ParamQuantKind:
# 如果参数名称以"embeddings.weight"结尾，返回ParamQuantKind.embedding_table表示该参数是嵌入表的权重。
if name.endswith("embeddings.weight"):
return ParamQuantKind.embedding_table
# 如果参数名称为"head.weight"，返回ParamQuantKind.final_fc_weight表示该参数是最后一个全连接层的权重。
elif name == "head.weight":
return ParamQuantKind.final_fc_weight
# 如果参数的维度为2且名称以".weight"结尾，返回ParamQuantKind.linear_weight表示该参数是线性层的权重。
elif param_info.ndim == 2 and name.endswith(".weight"):
return ParamQuantKind.linear_weight
else:
return ParamQuantKind.others
上面已经完成了RWKV模型的定义，接下来是定义几个相关的TIR函数并定义一个最终的TIR模型获取函数。这里对创建prefill和decode的create_func函数以及最终的TIR模型获取函数get_model进行解析：
# 函数接受一个relax.BlockBuilder对象bb、一个ParamManager对象param_manager、一个RWKVConfig对象config、
# 一个QuantizationScheme对象quant_scheme和一个字符串类型的函数名称func_name（默认为"prefill"或"decode"）
def create_func(
bb: relax.BlockBuilder,
param_manager: ParamManager,
config: RWKVConfig,
quant_scheme: QuantizationScheme,
func_name=Literal["prefill", "decode"],
):
# 如果函数名称不是"prefill"或"decode"，则抛出ValueError异常。
if func_name not in ["prefill", "decode"]:
raise ValueError(f"func_name must be 'prefill' or 'decode', got {func_name}")
# 根据函数名称确定序列的长度seq_len，如果函数名称为"decode"，则将序列长度设为1，否则设为tir.Var("n", "int64")。
seq_len = 1 if func_name == "decode" else tir.Var("n", "int64")
# 在BlockBuilder的function上下文中创建函数func_name。
with bb.function(func_name):
# 创建一个RWKVForCausalLM模型对象model。
model = RWKVForCausalLM(config)
# 调用param_manager的register_params方法注册模型参数。
param_manager.register_params(
model, func_name, quant_scheme, get_param_quant_kind
)
# 创建一个输入占位符input_ids，形状为(1, seq_len)，数据类型为"int32"。
input_ids = nn.Placeholder((1, seq_len), dtype="int32", name="input_ids")
# Placeholder for compatibility to LLAMA
# 创建一个占位符all_seq_len_shape，用于兼容LLAMA。
all_seq_len_shape = relax.Var("place_holder", R.Object())
# 创建一个变量state，其值为包含多个R.Object()的元组，长度为config.num_hidden_layers * 5。
state = relax.Var("state", R.Tuple([R.Object()] * config.num_hidden_layers * 5))
with bb.dataflow():
# 调用model的前向传播函数，将input_ids和state作为输入，得到输出logits和状态列表states。
logits, states = model(input_ids, state)
# 将input_ids、all_seq_len_shape、state和模型的参数列表作为参数列表params。
params = [
input_ids,
all_seq_len_shape,
state,
] + model.parameters()
# 使用bb.emit_output将(logits, relax.Tuple(states))作为输出。
gv = bb.emit_output((logits, relax.Tuple(states)))
# 使用bb.emit_func_output将输出和参数列表params作为函数的输出。
bb.emit_func_output(gv, params)
# 获取构建好的模块mod和global function变量gv。
mod = bb.get()
gv = mod.get_global_var(func_name)
# 根据函数名称更新函数的属性，包括输入数量和tir_var_upper_bound（如果函数名称为"prefill"）。
f = mod[gv].with_attr("num_input", 3)
if func_name == "prefill":
f = f.with_attr("tir_var_upper_bound", {"n": config.max_sequence_length})
bb.update_func(gv, f)
# 定义了一个名为get_model的函数，接受两个参数：args和hf_config。
def get_model(args, hf_config):
# 从args中获取模型名称、最大序列长度和模型数据类型。
model_name = args.model
max_seq_len = args.max_seq_len
dtype = args.quantization.model_dtype
# 检查模型名称是否以"rwkv-"开头，如果不是，则抛出ValueError异常。
if not model_name.lower().startswith("rwkv-"):
raise ValueError(f"Unsupported model name: {model_name}")
# 使用hf_config和dtype创建一个RWKVConfig配置对象config。
config = RWKVConfig(**hf_config, dtype=dtype)
# 如果指定了最大序列长度max_seq_len，则将config的max_sequence_length属性设置为max_seq_len。
if max_seq_len != -1:
config.max_sequence_length = max_seq_len
# 创建一个ParamManager对象param_manager用于管理模型参数，以及一个relax.BlockBuilder对象bb用于构建计算图。
param_manager = ParamManager()
bb = relax.BlockBuilder()
# 调用一系列create_func函数，向bb中添加计算图的构建指令，
# 包括"prefill"和"decode"两个函数、KV Cache函数、softmax函数和元数据函数。
create_func(bb, param_manager, config, args.quantization, "prefill")
create_func(bb, param_manager, config, args.quantization, "decode")
create_kv_cache_func(bb, config)
create_softmax_func(bb, config)
create_metadata_func(
bb,
model_name=model_name,
# RNN model do not have window size limit
max_window_size=-1,
stop_tokens=[0],
add_prefix_space=False,
)
create_kv_cache_reset_func(bb, config)
# 通过调用bb的get方法获取构建好的模块mod。
mod = bb.get()
# 如果args.build_model_only为True，则直接返回模块mod、参数管理器param_manager、None和配置config。
if args.build_model_only:
return mod, param_manager, None, config
# 定义一个名为f_convert_pname_fwd的函数，用于转换前向参数名称。
def f_convert_pname_fwd(pname: str) -> List[str]:
if (
"key_weight" in pname
or "value_weight" in pname
or "receptance_weight" in pname
or "output_weight" in pname
or "head_weight" in pname
):
return [pname.replace("_weight", ".weight")]
else:
return [pname]
# 定义一个名为f_convert_param_bkwd的函数，用于转换反向参数。
def f_convert_param_bkwd(torch_pname: str, torch_param):
# torch_param: numpy.ndarray
import numpy as np # pylint: disable=import-outside-toplevel
# rescale_every
if config.rescale_every > 0 and "blocks." in torch_pname:
# based-on the assumption that the layer id is the second element in torch_pname
layer_id = int(torch_pname.split(".")[2])
if (
"attention.output.weight" in torch_pname
or "feed_forward.value.weight" in torch_pname
):
torch_param = torch_param / (2 ** (layer_id // config.rescale_every))
# reshape
if "time_" in torch_pname:
torch_param = torch_param.squeeze()
# convert dtype
if "time_decay" in torch_pname: # need fp32 for this
return [(torch_pname, -np.exp(torch_param.astype("float32")))]
elif "time_first" in torch_pname:
return [(torch_pname, torch_param.astype("float32"))]
else:
return [(torch_pname, torch_param.astype(config.dtype))]
# 调用param_manager的set_param_loading_func方法，设置参数加载函数。
param_manager.set_param_loading_func(
args.model_path, args.use_safetensors, f_convert_pname_fwd, f_convert_param_bkwd
)
# 返回模块mod、参数管理器param_manager、长度为参数数量的None列表和配置config。
return mod, param_manager, [None] * len(param_manager.param_names), config
自此，我们基本就有了搭建RWKV模型的全部流程，说白了就是用TVM的Relax语言手动一对一的把PyTorch实现翻译过去。
0x3. Transform举例
在mlc-llm有一些图层的优化，在 https://github.com/BBuf/mlc-llm-code-analysis/tree/main/mlc_llm/transform 这个文件里面，我们对其中的一些优化Pass做一下解析。
0x3.1 rewrite attention
代码如下：
# 导入了TVM的relax模块中的一些函数和类，以及TVM的script模块中的relax别名。
from tvm.relax.dpl import PatternContext, is_const, is_op, rewrite_call, wildcard
from tvm.script import relax as R
# 定义了一个名为rewrite_attention的函数，接收一个参数f。
def rewrite_attention(f):
# 使用wildcard()创建了三个通配符，分别赋值给Q、K和V。
Q = wildcard()
K = wildcard()
V = wildcard()
# 使用is_op()函数创建了三个操作模式，分别对应Q、K和V的维度重排操作，并将结果分别赋值给Q_BNSH、K_BNSH和V_BNSH。
Q_BNSH = is_op("relax.permute_dims")(Q)
K_BNSH = is_op("relax.permute_dims")(K)
V_BNSH = is_op("relax.permute_dims")(V)
# 使用is_op()函数创建了一个操作模式，对应K_BNSH的维度重排操作，并将结果赋值给K_BNSH_T。
K_BNSH_T = is_op("relax.permute_dims")(K_BNSH)
# 使用is_op()函数创建了一系列操作模式，对应矩阵乘法、除法、最大值、最小值、softmax以及另一个矩阵乘法操作。
# 这些操作模式（Attention）根据之前定义的通配符和常数匹配不同的计算图节点。
matmul1 = is_op("relax.matmul")(Q_BNSH, K_BNSH_T)
divide = is_op("relax.divide")(matmul1, is_const())
max = is_op("relax.maximum")(divide, is_const())
min = is_op("relax.minimum")(max, wildcard())
softmax = is_op("relax.nn.softmax")(is_op("relax.astype")(min))
matmul2 = is_op("relax.matmul")(is_op("relax.astype")(softmax), V_BNSH)
# 使用is_op()函数创建了一个操作模式，对应matmul2的维度重排操作，并将结果赋值给pattern。
pattern = is_op("relax.permute_dims")(matmul2)
# 定义了一个名为callback的回调函数，接收两个参数_和matchings。
# 该回调函数使用R.nn.attention函数构建一个新的计算图节点，并使用matchings字典中的匹配结果来填充该节点的参数。
def callback(_, matchings):
return R.nn.attention(
matchings[Q], matchings[K], matchings[V], causal_mask="BottomRight"
)
# 使用rewrite_call函数将pattern、callback和输入的计算图f传递给它，以便在计算图中应用模式匹配和重写。
# 最后，将重写后的计算图返回。
return rewrite_call(pattern, callback, f)
虽然没有完全看懂这里的操作比如max和min的含义，但是从后面的callback_可以猜测出这里的Pass就是把打散的Self Attention模块融合为一个relax.nn.attention操作。在cuda后端，如果支持了cutlass，那么relax.nn.attention操作就对应了Flash Attention。
0x3.2 Transpose MatMul
代码实现解析如下：
# 这段代码定义了一个名为TransposeMatmulCodeGenerator的类，该类继承自relax.PyExprMutator。
# 通过@relax.expr_functor.mutator装饰器将该类声明为一个表达式重写器。
@relax.expr_functor.mutator
class TransposeMatmulCodeGenerator(relax.PyExprMutator):
def __init__(self, mod):
super().__init__(mod)
@staticmethod
def pattern():
# 定义了静态方法pattern()，该方法返回一个描述模式的元组。
# 通过使用通配符(wildcard())和操作模式(is_op())来匹配计算图中的特定模式。
# 在这个例子中，模式匹配了一个矩阵乘法操作中矩阵w的维度重排操作，并将匹配的结果保存在字典annotations中。
w = wildcard()
x = wildcard()
wT = is_op("relax.permute_dims")(w)
o = is_op("relax.matmul")(x, wT)
annotations = {"o": o, "w": w, "x": x, "wT": wT}
# 定义了内部函数_check()，用于检查模式匹配的结果是否满足特定的条件。
# 在这个例子中，检查了维度重排操作的维度数和轴的顺序是否正确。
def _check(context: relax.transform.PatternCheckContext) -> bool:
transpose_call = context.annotated_expr["wT"]
ndim = transpose_call.args[0].struct_info.ndim
if ndim == -1:
return False
if ndim == 2 and transpose_call.attrs.axes is None:
return True
axes = list(range(ndim))
axes[-1], axes[-2] = axes[-2], axes[-1]
return list(transpose_call.attrs.axes) == axes
# 将匹配的计算图节点、注解和检查函数作为元组返回。
return o, annotations, _check
# 重写了父类的visit_call_()方法，用于处理特定类型的计算图节点。
def visit_call_(self, call: relax.Call) -> relax.Expr:
# 定义了一个变量out_dtype，用于保存输出的数据类型。
out_dtype = None
# 定义了一个内部函数te_transposed_matmul()，该函数实现了矩阵乘法的计算逻辑。
def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
nonlocal out_dtype
# 将输入张量 a 和 b 的形状转换为列表形式，分别保存在变量 a_shape 和 b_shape 中。
a_shape = list(a.shape)
b_shape = list(b.shape)
# 定义了两个布尔变量 a_prepended 和 b_appended，用于标记是否在相应的形状的前面或后面添加了维度。
a_prepended = False
b_appended = False
# 如果输入张量 a 的形状为一维，则在其前面添加一个维度，将其形状修改为 (1, original_shape)。
# 同样地，如果输入张量 b 的形状为一维，则在其后面添加一个维度，将其形状修改为 (original_shape, 1)。
if len(a_shape) == 1:
a_prepended = True
a_shape.insert(0, 1)
if len(b_shape) == 1:
b_appended = True
b_shape.append(1)
# 比较 a_shape 和 b_shape 的长度，将结果保存在布尔变量 is_a_larger 中。
# offset 表示两个形状长度之差，用于后续处理。
is_a_larger = len(a_shape) > len(b_shape)
offset = (
len(a_shape) - len(b_shape)
if is_a_larger
else len(b_shape) - len(a_shape)
)
# 创建两个 relax.Var 对象 a_relax 和 bT_relax，用于表示张量 a 和转置后的张量 bT 的结构信息。
# a_relax 的形状和 a 的形状相同，bT_relax 的形状是 b 的形状经过维度互换后的结果。
a_relax = relax.Var("a", relax.TensorStructInfo(a.shape))
bT_shape = list(b.shape)
bT_shape[-1], bT_shape[-2] = bT_shape[-2], bT_shape[-1]
bT_relax = relax.Var("b", relax.TensorStructInfo(bT_shape))
# 使用 relax.op.matmul() 方法对 a_relax 和 bT_relax 进行矩阵乘法运算。
# 然后，通过 self.builder_.normalize() 方法对结果进行归一化处理，并获取最终的输出形状。
output_shape = self.builder_.normalize(
relax.op.matmul(a_relax, bT_relax)
).struct_info.shape
# 该函数接受可变数量的空间索引参数 idx_spatial，
def matmul_compute(*idx_spatial):
# 并定义了一个名为 k 的规约轴（reduce axis），其范围为 0 到 a_shape[-1]。
k = te.reduce_axis((0, a_shape[-1]), name="k")
# 定义了一个名为 multiply_compute 的内部函数，用于计算乘法操作时的索引。
def multiply_compute(idx_reduce):
a_indices = []
b_indices = []
# 根据 is_a_larger 的值，将 idx_spatial 中的索引分配给 a_indices 或 b_indices，用于处理形状长度差异的维度。
for i in range(offset):
if is_a_larger:
a_indices.append(idx_spatial[i])
else:
b_indices.append(idx_spatial[i])
for i in range(
offset, len(output_shape) - (2 - a_prepended - b_appended)
):
# 根据维度的相等性，将适当的索引添加到 a_indices 和 b_indices 中。
# 如果维度不相等或无法确定是否相等，则将索引设为 0 或保持不变。
a_dim = a_shape[i if is_a_larger else i - offset]
b_dim = b_shape[i if not is_a_larger else i - offset]
dim_equal = a_dim == b_dim
if not isinstance(dim_equal, tir.IntImm) or dim_equal == 0:
a_dim_is_one = isinstance(a_dim, tir.IntImm) and a_dim == 1
b_dim_is_one = isinstance(b_dim, tir.IntImm) and b_dim == 1
a_indices.append(0 if a_dim_is_one else idx_spatial[i])
b_indices.append(0 if b_dim_is_one else idx_spatial[i])
else:
a_indices.append(idx_spatial[i])
b_indices.append(idx_spatial[i])
# 在乘法操作的索引中添加规约轴 idx_reduce，并根据 a_prepended 和 b_appended 的值，
# 将适当的索引添加到 a_indices 和 b_indices 中。
if not a_prepended:
a_indices.append(idx_spatial[-2 + b_appended])
a_indices.append(idx_reduce)
if not b_appended:
b_indices.append(idx_spatial[-1])
b_indices.append(idx_reduce)
# 根据 out_dtype 的值，选择是否进行数据类型转换，并返回乘法操作的结果。
dtype = out_dtype
if dtype != "":
return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
return a(*a_indices) * b(*b_indices)
# 在缩减轴 k 上对 multiply_compute 的结果进行求和操作。
return te.sum(multiply_compute(k), axis=k)
# 使用 te.compute() 函数计算最终的输出，其中使用一个 lambda 函数将输入索引传递给 matmul_compute 函数，
# 并将结果命名为 "NT_matmul"。整个计算过程将根据 output_shape 进行执行。
return te.compute(
output_shape,
lambda *idx: matmul_compute(*idx), # pylint: disable=unnecessary-lambda
name="NT_matmul",
)
# 首先，检查函数调用的操作符 call.op 是否是 relax.GlobalVar 类型。如果是，获取与该操作符对应的函数对象，
# 并检查函数的属性中是否包含键 "Composite"，且其值为 "transpose_matmul_fuse"。
if isinstance(call.op, relax.GlobalVar):
function = self.builder_.get()[call.op]
if (
"Composite" in function.attrs
and function.attrs["Composite"] == "transpose_matmul_fuse"
):
# 将函数的返回类型 function.ret_struct_info.dtype 赋值给变量 out_dtype
out_dtype = function.ret_struct_info.dtype
# 然后调用 self.builder_.call_te() 方法，传递 te_transposed_matmul 函数作为参数，
# 以及调用的参数 call.args[1] 和 call.args[0]，并指定 primfunc_name_hint 为 "NT_matmul"。
return self.builder_.call_te(
te_transposed_matmul,
call.args[1],
call.args[0],
primfunc_name_hint="NT_matmul",
)
return super().visit_call_(call)
# 使用 @tvm.transform.module_pass 装饰器定义了一个名为 FuseTransposeMatmul 的类，
# 并指定了优化级别 opt_level=0 和 pass 的名称为 "FuseTransposeMatmul"。
@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
class FuseTransposeMatmul:
# 定义了 transform_module 方法，接受一个名为 mod 的 IRModule 对象和
# tvm.transform.PassContext 对象作为参数，并返回一个 IRModule 对象。
def transform_module(
self, mod: IRModule, ctx: tvm.transform.PassContext
) -> IRModule:
# 通过调用 relax.transform.FuseOpsByPattern 并传递一个包含单个模式元组的列表，
# 对模块 mod 进行融合的转置矩阵乘法操作。
mod = relax.transform.FuseOpsByPattern(
[("transpose_matmul_fuse", *TransposeMatmulCodeGenerator.pattern())]
)(mod)
# 创建一个名为 transpose_matmul_codegen 的 TransposeMatmulCodeGenerator 对象，
# 并对模块中的每个函数进行遍历。如果函数是 relax.Function 类型，则调用 transpose_matmul_codegen.visit_expr
# 方法对函数进行转置矩阵乘法代码生成，并通过 transpose_matmul_codegen.builder_.update_func 方法更新函数。
transpose_matmul_codegen = TransposeMatmulCodeGenerator(mod)
for gv in mod.functions:
func = mod[gv]
if not isinstance(func, relax.Function):
continue
func = transpose_matmul_codegen.visit_expr(func)
transpose_matmul_codegen.builder_.update_func(gv, func)
# 返回转置矩阵乘法代码生成器的 builder 对象中的模块。
return transpose_matmul_codegen.builder_.get()
这个Pass将Transpose算子和一个MatMul算子替换为一个TE表达式的实现来达到融合算子的目的。
除了上面2种Pass，MLC-LLM还有不少的图变换Pass，这篇文章就不一一去解析了，大多数优化的目的都是匹配某种Pattern然后用更优秀的算子去完成计算。
量化策略这一块就不在这篇文章解析了。
0x4. MLC-LLM优缺点个人评价和期待
0x4.1 优点
- Tune Free。mlc-llm不需要用TVM的AutoTVM/Ansor等等程序去执行算子搜索过程，对跨平台部署是比原始的TVM搭建的模型更清真的。
- TIR的语法很大程度靠近了PyTorch的API，使得用户在模型搭建部分不会很困难。
- 文档写得不错，跟随教程基本可以完成大多数平台的模型部署，并且单Batch下的吞吐和延迟表现都是不错的。
0x4.2 缺点
- 不支持从onnx或者huggingface模型直接转换出TIR，手工实现模型的时候需要相当多的先验知识，比如在上面的RWKV模型中如果有自定义的cuda kernel，那么这个模型的实现可能只能全权委托给mlc-ai社区的核心开发人员了。
- KV Cache开的是max_sequence_length这么长，显然会有显存的浪费，Serving的时候极限情况下可以服务的用户数量应该比VLLM/TGI等要小？
- CUDA后端Decoding的Attention我看起来好像还是会用Flash Attention？也许是我看错了，这条暂时存疑。
- 在RWKV模型实现里，看到Batch维度写死为1了，应该不支持动态Batch？这样对于启真实服务来说会有一些限制。
0x4.3 期待
- 如果短期内能让一个对TVM只有轻度依赖的社区开发者新增一个新的模型。
- 如果模型存在自定义CUDA Kernel，需要一个详细的教程来指引。
- 模型逐层打印来debug精度缺一个教程。
- Paged Attention类似策略的引入。
- 动态Batch的支持。
暂时就想到这些，欢迎斧正。