add triton based Softmax and unit tests

[lessw2020]

Summary:
This PR is to add a Triton based softmax, which supports both fwd and bwd as well as multi-dim inputs.
This kernel outperforms the PT native softmax in the fwd pass and thus might be favored esp. for inference situations.

It's also a first opportunity to consider expanding the TorchMM offerings to include triton kernels.
Thus, I have listed this under a separate Triton directory in order to separate it from the rest of the codebase which is PyTorch base.

Test plan:
Tested it in end to end training (short run).
Added a number of unit tests which compare it directly with the torch.softmax and tests both forward and backward passes, as well as 2d and multi-dim inputs and fp32 and bf16 dtypes.

[rohan-varma]

Curious, when would users have to author this sort of performant triton-based softmax, versus getting these perf wins ootb with torch.compile?