Method forward

tensorrt_llm/layers/mlp.py:231–264

(self,
                hidden_states,
                lora_layer_params=None,
                all_reduce_params: Optional[AllReduceParams] = None)

                                 gather_output=False)

    def forward(self,
                hidden_states,
                lora_layer_params=None,
                all_reduce_params: Optional[AllReduceParams] = None):
        if lora_layer_params is not None:
            assert lora_layer_params.get_runtime_params(
                0, "mlp_gate_up"
            ) is None, f"LoRA module 'mlp_gate_up' is not supported in {self}"

        mlp_fc_lora_params = None
        if lora_layer_params is not None:
            mlp_fc_lora_params = lora_layer_params.get_runtime_params(
                0, "mlp_h_to_4h")

        mlp_gate_lora_params = None
        if lora_layer_params is not None:
            mlp_gate_lora_params = lora_layer_params.get_runtime_params(
                0, "mlp_gate")

        mlp_proj_lora_params = None
        if lora_layer_params is not None:
            mlp_proj_lora_params = lora_layer_params.get_runtime_params(
                0, "mlp_4h_to_h")

        inter = self.fc(hidden_states, mlp_fc_lora_params)
        inter = ACT2FN[self.hidden_act](inter)
        gate = self.gate(hidden_states, mlp_gate_lora_params)
        intermediate = inter * gate
        if self.inner_layernorm is not None:
            intermediate = self.inner_layernorm(intermediate)
        output = self.proj(intermediate,
                           lora_runtime_params=mlp_proj_lora_params,
                           all_reduce_params=all_reduce_params)
        return output


class FusedGatedMLP(Module):