Extract forward/backward core computation bits outside of torch autograd context for easy reuse

[qingquansong]

Summary

Extract forward/backward core computation bits outside of torch autograd context for easy reuse. This is beneficial for lightning thunder integration and the reuse of kernel in other context.

Doubled checked the speed and memory usage, within variance range, no degradation

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[ByronHsu]

cc @t-vi LGTM