Add backward kernel for exp

[aditvenk]

Test Plan:
Adding new unit tests in test_examples.py
Testing with run.py and exp.py

[oulgen]

Please add this to run.py, you can grep for rms_norm again for backward implementation, and share perf/accuracy results with tritonbench.

@Sibylau can give pointers on how to do it

gist is here: https://docs.google.com/document/d/1BiQaFJRBufzcLNPSMhVKuC9gv-KEFpN4B49CqOjZC74/edit?usp=sharing

[aditvenk]

Please add this to run.py, you can grep for rms_norm again for backward implementation, and share perf/accuracy results with tritonbench.

@Sibylau can give pointers on how to do it

gist is here: https://docs.google.com/document/d/1BiQaFJRBufzcLNPSMhVKuC9gv-KEFpN4B49CqOjZC74/edit?usp=sharing

From what I see, tritonbench does not have a bwd operator for this op (the same is true for some other ops too)
I will check-in with @Sibylau how to proceed here.

[yf225]

Please add this to run.py, you can grep for rms_norm again for backward implementation, and share perf/accuracy results with tritonbench.
@Sibylau can give pointers on how to do it
gist is here: https://docs.google.com/document/d/1BiQaFJRBufzcLNPSMhVKuC9gv-KEFpN4B49CqOjZC74/edit?usp=sharing

From what I see, tritonbench does not have a bwd operator for this op (the same is true for some other ops too) I will check-in with @Sibylau how to proceed here.

Yeah let's probably also add backward support for this operator to tritonbench (we can look at get_bwd_fn in tritonbench layer_norm and rms_norm for how to do this; also need to make sure the inputs have proper requires_grad=True too).

[aditvenk]

Please add this to run.py, you can grep for rms_norm again for backward implementation, and share perf/accuracy results with tritonbench.
@Sibylau can give pointers on how to do it
gist is here: https://docs.google.com/document/d/1BiQaFJRBufzcLNPSMhVKuC9gv-KEFpN4B49CqOjZC74/edit?usp=sharing

From what I see, tritonbench does not have a bwd operator for this op (the same is true for some other ops too) I will check-in with @Sibylau how to proceed here.

Yeah let's probably also add backward support for this operator to tritonbench (we can look at get_bwd_fn in tritonbench layer_norm and rms_norm for how to do this; also need to make sure the inputs have proper requires_grad=True too).

tritonbench PR: meta-pytorch/tritonbench#501