add fused transpose and non-transpose kernel and use it for grad output

[danielvegamyhre]

Stack Summary
The following stack of PRs completes a float8 training prototype with performance that slightly beats the production Float8Linear + torch.compile approach.

Changes in this PR
This PR implements a new kernel which reads in a row major high precision tensor, and writes 2 outputs: a fp8 row major output, and a transposed fp8 row major output. This is useful for the backward pass where we need to convert the grad_output to both of these formats.

Next steps:

1. Add support for activation checkpointing
2. Add test verifying this prototype is compatible with FSDP
3. Implement usage of this prototype in torchtitan and benchmark results

Test Plan

1. pytest kernels/ - kernel specific unit tests are passing
2. pytest test/ - e2e training test is passing

Performance Benchmarking

Performance benchmarks show this implementation is beating torch.compile by 1.72-4.45% depending on the input tensor size:

input_shape    kernel_algo                 high_precision_dtype      eager_time    compiled_time    float8nocompile
-------------  --------------------------  ----------------------  ------------  ---------------  -----------------
(16, 4096)     KernelAlgorithm.ATOMIC_MAX  torch.bfloat16               649.218          394.725            386.469
(256, 4096)    KernelAlgorithm.ATOMIC_MAX  torch.bfloat16               685.783          420.743            408.137
(4096, 4096)   KernelAlgorithm.ATOMIC_MAX  torch.bfloat16              1829.13          1053.64             977.858
(65536, 4096)  KernelAlgorithm.ATOMIC_MAX  torch.bfloat16             21554.2          12369.7            10813.3
(16, 4096)     KernelAlgorithm.REDUCTION   torch.bfloat16               650.026          394.951            696.221
(256, 4096)    KernelAlgorithm.REDUCTION   torch.bfloat16               684.865          421.144            729.459
(4096, 4096)   KernelAlgorithm.REDUCTION   torch.bfloat16              1826.42          1050.85            1596.12
(65536, 4096)  KernelAlgorithm.REDUCTION   torch.bfloat16             21584.7          12347.2            17290

[danielvegamyhre]

Stack from ghstack (oldest at bottom):

• float8nocompile: add e2e fsdp test #1523
• support for activation checkpointing in float8nocompile #1517
• fix bug in tl.store mask for kernel _to_fp8_row_major_t_and_non_t #1516
• add handling for batch dim in float8nocompile #1512
• -> add fused transpose and non-transpose kernel and use it for grad output #1497

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1497

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 01eedbf with merge base eb49333 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.