[SLM] Fuse Add and RMSNorm

[jinhongyii]

This PR adds a fusion pass that applies to "binary add" and "RMSNorm".
This is a temporary workaround that allows us to fuse "add" into RMSNorm
once it is not fused into GEMM epilogue.

[jinhongyii]

cc: @MasterJH5574

[junrushao]

I suppose the fused add-rmsnorm operator could be expressed by TE and scheduled by Dlight easily

[jinhongyii]

No it's not that easy. The tricky point is that add_rmsnorm function has 2 outputs: both results of add and rms_norm, which makes compute_inline/compute_at/reverse_compute_at all fail in such case. I have to write scheduled TIR to work around.

[junrushao]

Got it. Thanks for the elaboration!

[vinx13]

What's the implication on performance? If the manual schedule can't generalize to all cases we can try supporting such pattern in cublas fusion

[jinhongyii]

I heard from @MasterJH5574 that cublas fusion cannot successfully fuse matmul and divide_add now, so I create this small pass to unblock our effort on mlc serve perf profiling. In the long term, surely this can be replaced by better cublas fusion, but it doesn't hurt to work as a fallback or as a target to compare.

[junrushao]

Whenever we use TIR, I know they are auto generated ones, but lets work to make sure they are human readable so that they establish positive examples demonstrating “TIR is actually great”

[jinhongyii]

Thanks @junrushao for pointing out. I accidentally used black to format this TIR and I will regenerate this TIR.

[junrushao]

I think the PR is ready to merge in terms of code quality and correctness. I believe 1) @vinx13 has some further comments on its performance implication and generalizability, and 2) I have some concern about the prefill kernel which assumes batch_size = 1. Anyways, it's a good start.