Another MLP implementation along with multiplier support

[Tcc0403]

🚀 The feature, motivation and pitch

Currently, LigerMLP modules only fuse swiglu/geglu computations together and leave matmuls untouched. These elementwise operations (including multiplier support #936 ) could be easily fused into matmul's epilogues. We can investigate the performance of this approach and see if we should adopt it.

TL;DR

Instead of

gate_states = self.gate_proj(x)
up_states = self.up_proj(x)
intermidiate_states = LigerSiLUMulFunction.apply(gate_states , up_states)
return self.down_proj(intermidiate_states)

There are some other approaches worth exploring:

1. fuse activations (and multiplier) into gate_proj(x)

up_states = self.up_proj(x)
intermidiate_states = LigerFusedLinearActMultiplierFunction.apply(x, self.gate_proj.weight, gate_multiplier, up_states)
return self.down_proj(intermidiate_states)

2. stack gate and up projections then put it into activation functions

gate_up_states = self.gate_up_proj(x)
intermidiate_states = LigerSplitStatesActMultiplierFunction.apply(
	gate_up_states, 
	config.hidden_act, 
	gate_multiplier, 
	up_states
)
return self.down_proj(intermidiate_states)

3. dual gemm with activations (and multiplier)

intermidiate_states = LigerDualGemmActMulFuncion.apply(
	x, 
	self.gate_proj.weight, 
	self.up_proj.weight, 
	config.hidden_act, 
	gate_multiplier,
)
return self.down_proj(intermidiate_states)

Alternatives

No response

Additional context

No response

[0xtoward]

Hi! I'm about to work on 936 and 937. My plan is to:
PR1:

1. Add LigerFalconSwiGLUMLP (transformers/swiglu.py) for both multipliers. (or edit current SwiGLUMLP instead?)

2. Support Falcon H1: Enable the MLP patch in apply_liger_kernel_to_falcon_h1 and remove the current warning.

PR2:

1. Extend the LigerSiLUMulFunction (ops/swiglu.py) for fused gate_multiplier.

Does this design fit Liger? Thanks!

[Tcc0403]

Yes. Feel free to ping me for any questions or requesting reviews!

[MyDarapy]

#take Hi @Tcc0403, I’d like to investigate this and work on this. I checked current main and confirmed that #1201 only adds multiplier support but leaves the projection GEMMs unfused. I propose starting with a benchmark-first implementation that fuses gate projection + SiLU + multiplier + multiplication with the precomputed up projection.