slapo.op.mlp

MLP module using with fused kernels.

Classes:

FusedMLP(hidden_size, intermediate_size, ...)

A wrapper MLP to make use of fused bias+gelu and bias+dropout.

class slapo.op.mlp.FusedMLP(hidden_size, intermediate_size, orig_act, resid_pdrop, use_torchscript=False)

A wrapper MLP to make use of fused bias+gelu and bias+dropout. Note that both linear modules in this MLP have bias, so users should not replace the original MLP with this module if the original MLP does not have bias.

Parameters

• hidden_size (int) – The hidden size of the input.

• intermediate_size (int) – The intermediate size of the MLP.

• orig_act (str) – The original activation function in string.

• resid_pdrop (float) – The dropout probability for the residual connection.

• use_torchscript (bool) – Whether to use torchscript or memory_efficient_fusion.

Methods:

forward(hidden_states)

Defines the computation performed at every call.

forward(hidden_states)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.