[mxfp8 moe training] use dim1 cast cuda kernel in bwd

[danielvegamyhre]

Stacked PRs:

• ->[mxfp8 moe training] use dim1 cast cuda kernel in bwd #2897

---

[mxfp8 moe training] use dim1 cast cuda kernel in bwd

This way we can eliminate the .contiguous() call, as well as use the dim1 cast CUDA kernel which is the fastest option (vs inductor codgen or triton).

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm CI/CD workflows failing due to : download from https://api.github.com/repos/pytorch/pytorch timed out.

❌ 2 Cancelled Jobs, 4 Pending

As of commit e64d4b5 with merge base 83a20c7 ():

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  ##[error]The operation was canceled.

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

[vkuzo]

nit: move to utils file if you want to reuse it

[danielvegamyhre]

done