Investigate triton

[epwalsh]

I've been messing around with triton to see if it makes sense to start replacing some of our components with a triton implementation. So far preliminary results look good.

I have been working on a triton version of LayerNorm both with and without the element-wise affine transform. These are the benchmarking results for a batch of 4096 tokens and d_model of 4096 (representative of a typical microbatch with our medium model) or 8192 (our large model), on an A100 GPU.

The units are in GBPS (throughput), so larger is better.

layer-norm-with-affine-forward:
          N       Triton        Torch
1    4096.0  1129.931006   936.228546
2    8192.0  1379.705219   949.797080

layer-norm-with-affine-backward:
          N      Triton       Torch
1    4096.0  309.132070  479.531698
2    8192.0  632.180043  491.520012

layer-norm-no-affine-forward:
          N       Triton        Torch
1    4096.0  1092.266694   992.969689
2    8192.0  1409.376308   949.797080

layer-norm-no-affine-backward:
          N       Triton       Torch
1    4096.0  1156.517652  750.412251
2    8192.0  1524.093109  712.347810

[epwalsh]

Marking as blocked again because it doesn't appear to work properly on AMD GPUs. See #260.

[dumitrac]

Marking the items prior to Feb 29th as "closed".