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

Code before:
        q, k, v = op_ext.llama_rope(
            qkv, total_seq_len, self.rope_theta, self.num_attention_heads, self.num_key_value_heads
        )
        # 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(total_seq_len)
        v = self.v_cache.view(total_seq_len)
        # 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)


ACT2FN = {
    "gelu": partial(nn.gelu, approximate=False),
    "relu": nn.relu,
    "silu": nn.silu,
    "swish": nn.silu,
    "gelu_new": partial(nn.gelu, approximate=True),
}


class QWen2MLP(nn.Module):
    def __init__(self, config: QWen2Config):
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_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.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x: Tensor):
        gate_oup = self.gate_proj(x)
        act_oup = self.act_fn(gate_oup)
        up_oup = self.up_proj(x)
        down_oup = self.down_proj(act_oup * up_oup)
        return down_oup


class QWen2DecoderLayer(nn.Module):
    def __init__(self, config: QWen2Config):
        self.self_attn = QWen2Attention(config)
        self.mlp = QWen2MLP(config)
        self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False)
        self.post_attention_layernorm = nn.RMSNorm(
            config.hidden_size, -1, config.rms_norm_eps, bias=False
        )

    def forward(self, hidden_states: Tensor, attention_mask: Tensor, total_seq_len: tir.Var):
        out = self.input_layernorm(hidden_states)
        out = self.self_attn(out, attention_mask, total_seq_len)
        hidden_states = out + hidden_states

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:
        q, k, v = op_ext.llama_rope(
            qkv, total_seq_len, self.rope_theta, self.num_attention_heads, self.num_key_value_heads
        )
        # 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(total_seq_len)
        v = self.v_cache.view(total_seq_len)
        # 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)


ACT2FN = {
    "gelu": partial(nn.gelu, approximate=False),
    "relu": nn.relu,
    "silu": nn.silu,
    "swish": nn.silu,
    "gelu_new": partial(nn.gelu, approximate=True),
}


class QWen2MLP(nn.Module):
    def __init__(self, config: QWen2Config):
        self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x: Tensor):
        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)


class QWen2DecoderLayer(nn.Module):
    def __init__(self, config: QWen2Config):
        self.self_attn = QWen2Attention(config)
        self.mlp = QWen2MLP(config)
        self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False)
        self.post_attention_layernorm = nn.RMSNorm(
            config.hidden_size, -1, config.rms_norm_eps, bias=False
        )

    def forward(self, hidden_states: Tensor, attention_mask: Tensor, total_seq_len: tir.Var):
        out = self.input_layernorm(hidden_states)
        out = self.self_attn(out, attention_mask, total_seq_len)
        hidden_states = out + hidden_states

        out = self.post_attention_layernorm(hidden_states)
        out = self.mlp(out)