Repository: mlc-ai/mlc-llm
File: python/mlc_chat/model/baichuan/baichuan_model.py (Python)

Code before:
        hidden_states: Tensor,
        attention_mask: Tensor,
        total_seq_len: tir.Var,
    ):
        d, h, t = self.head_dim, self.num_heads, total_seq_len
        b, s, _ = hidden_states.shape
        assert b == 1, "Only support batch size 1 at this moment."
        # Step 1. QKV Projection
        qkv = self.W_pack(hidden_states)
        qkv = op.reshape(qkv, (b, s, 3 * h, d))
        # Step 2. Apply QK rotary embedding
        q, k, v = op_ext.llama_rope(qkv, t, 10000, h, h)
        # Step 3. Query and update KVCache
        self.k_cache.append(op.squeeze(k, axis=0))
        self.v_cache.append(op.squeeze(v, axis=0))
        k = self.k_cache.view(t)
        v = self.v_cache.view(t)
        # Step 4. Compute softmax(Q @ K^T / sqrt(d)) @ V
        output = op_ext.attention(q, k, v, casual_mask=attention_mask)
        # Step 5. Apply output projection
        return self.o_proj(output)


class BaichuanMLP(nn.Module):
    def __init__(self, config: BaichuanConfig):
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)

    def forward(self, x):
        return self.down_proj(op.silu(self.gate_proj(x)) * self.up_proj(x))


class BaichuanDecoderLayer(nn.Module):
    def __init__(self, config: BaichuanConfig):
        norm_eps = config.rms_norm_eps
        self.self_attn = BaichuanAttention(config=config)
        self.mlp = BaichuanMLP(config)
        self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, norm_eps, bias=False)
        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, norm_eps, bias=False)

    def forward(self, hidden_states: Tensor, attention_mask: Tensor, total_seq_len: tir.Var):
        out = self.self_attn(self.input_layernorm(hidden_states), attention_mask, total_seq_len)
        hidden_states = out + hidden_states
        out = self.mlp(self.post_attention_layernorm(hidden_states))
        hidden_states = out + hidden_states
        return hidden_states


class BaichuanModel(nn.Module):
    def __init__(self, config: BaichuanConfig):

Review comment (suggestion):
Here, we can combine the gate_proj and up_proj matmul similar to what we did in combing QKV in attention layer.
```
self.gate_up_proj = nn.Linear(
    in_features=config.hidden_size,
    out_features=2 * config.intermediate_size,
    bias=False,
)
self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False)
```

And in `forward`,
```
concat_x1_x2 = self.gate_up_proj(x)
x1, x2 = op.split(concat_x1_x2, 2, axis=-1)
return self.down_proj(self.act_fn(x1) * x2)
```

Also need to transform the parameters in the loader, similar to this one: https://github.com/mlc-ai/mlc-llm/blob/main/python/mlc_chat/model/llama/llama_loader.py#L65

Code after:

    def forward(  # pylint: disable=too-many-locals
        self,
        hidden_states: Tensor,
        attention_mask: Tensor,
        total_seq_len: tir.Var,
    ):
        d, h, t = self.head_dim, self.num_heads, total_seq_len
        b, s, _ = hidden_states.shape
        assert b == 1, "Only support batch size 1 at this moment."
        # Step 1. QKV Projection
        qkv = self.W_pack(hidden_states)
        qkv = op.reshape(qkv, (b, s, 3 * h, d))
        # Step 2. Apply QK rotary embedding
        q, k, v = op_ext.llama_rope(qkv, t, 10000, h, h)
        # Step 3. Query and update KVCache
        self.k_cache.append(op.squeeze(k, axis=0))
        self.v_cache.append(op.squeeze(v, axis=0))
        k = self.k_cache.view(t)
        v = self.v_cache.view(t)
        # Step 4. Compute softmax(Q @ K^T / sqrt(d)) @ V
        output = op_ext.attention(q, k, v, casual_mask=attention_mask)
        # Step 5. Apply output projection
        return self.o_proj(output)


class BaichuanMLP(nn.Module):
    def __init__(self, config: BaichuanConfig):
        self.gate_up_proj = nn.Linear(
            in_features=config.hidden_size,
            out_features=2 * config.intermediate_size,
            bias=False,
        )
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        concat_x1_x2 = self.gate_up_proj(x)
        x1, x2 = op.split(concat_x1_x2, 2, axis=-1)
        return self.down_proj(op.silu(x1) * x2)


class BaichuanDecoderLayer(nn.Module):
    def __init__(self, config: BaichuanConfig):
        norm_eps = config.rms_norm_eps
        self.self_attn = BaichuanAttention(config=config)
        self.mlp = BaichuanMLP(config)
        self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, norm_eps, bias=False)
        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, norm_eps, bias=False)

    def forward(self, hidden_states: Tensor, attention_mask: Tensor, total_seq_len: tir.Var):
        out = self.self_attn(self.input_layernorm(hidden_states), attention_mask, total_seq_len)