Combining residual block and layer norm fusion

[masahi]

Hi, I'm looking at optimization for stable diffusion. One very common pattern is gemm -> residual add -> layer norm. In CUTLASS, both residual block and layer norm fusion are supported in isolation (via

cutlass/include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h

Line 70 in 277bd6e

struct GemmWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false> {

and

cutlass/include/cutlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h

Line 63 in 66d9cdd

struct GemmLayernormMainloopFusion {

), but since the two kernels don't seem to "compose" each other, I don't see a way to realize residual + layer norm fusion simultaneously.

Am I missing something? This should be a very common problem people encounter in practice... so I hope there is a good solution. Or is this also what the V3 API can solve?

@hwu36 @thakkarV

[masahi]

Looking more closely at how layer norm fusion is supported, I'm assuming that layer norm fusion in cutlass is intended to be used for gemm -> layer norm -> gemm fusion. Since the residual add happens before layer norm, I think what I need to do is to replace the first gemm implementation https://github.com/NVIDIA/cutlass/blob/master/examples/37_gemm_layernorm_gemm_fusion/gemm_with_epilogue_visitor.h with the one that supports residual addition.

The situation seems much better now.

[masahi]

Actually it looks even simpler: I think I just need to modify https://github.com/NVIDIA/cutlass/blob/master/examples/37_gemm_layernorm_gemm_fusion/gemm_with_layernorm.h#L525 to support residual addition? This is the code for first GEMM epilogue.

[jackkosaian]

Hi, @masahi. Yes, I believe you can change ElementFunctorOp here to be LinearCombinationResidualBlock (defined here).

You would also need to change the Arguments structures (e.g., here) so that you could pass in the residual tensor.

[MingZwhy]

Hi, @jackkosaian , If I change ElementFunctorOp into LinearCombinationResidualBlock, will the EpilogueFunctorOp of the second gemm also be changed? If so, can I set the EpilogueFunctorOp for two gemm splitly? For example, set the EpilogueFunctorOp of the first gemm to be residual and the EpilogueFunctorOp of the second gemm to be relu. Thanks for your advise!

[jackkosaian]

@MingZwhy , you should be able to do this. You'll just need to do some plumbing so that the EpilogueFunctorOp passed to GemmMainloopFusion (here) uses ReLU.