Optimize the Trompt example to reduce training time by ~30%

[akihironitta]

Summary

Reduces training time by 29% by:

• removing unnecessary stream and D2H synchronisations (part of Remove unnecessary CUDA synchronisations #453)
• replacing einsum with matmul
• making data copied async wrt host
• using TF32 for matmuls

Benchmark results

Benchmarked the change using jannis dataset on g6.4xlarge that has a L4 GPU. To reproduce, run:

$ python examples/trompt.py --dataset jannis
$ python examples/trompt.py --dataset jannis --compile

Average time per training step:

	master (A)	this PR (B)	B/A
eager	742.79 ms	741.22 ms	0.99x
compile	262.24 ms	186.80 ms	0.71x

Average time per evaluation step:

	master (A)	this PR (B)	B/A
eager	303.76 ms	302.73 ms	1.00x
compile	58.75 ms	56.27 ms	0.96x

---

Happy holidays! 🎄🧹