Add mx_fp8_bf16 kernel

[drisspg]

Stacked PRs:

• Add mx_fp4_kernel #1661
• ->Add mx_fp8_bf16 kernel #1637
• Update to cutlass 3.8 | wait for tag to land #1634

---

Add mx_fp8_bf16 kernel

Will flesh out more but this moves over the kernel from here: https://github.com/drisspg/driss_torch/blob/2813322f0b0f9a0f0fc8d382090ad0aaecf3468a/src/mx_fp8_bf16.cu#L162

This does fp8xfp8 w/ E8m0 scales and group_size hard coded to 32. The format for the scales is the same as that for cublasLT. I have created a pytorch function that converts the [n_rows, n_cols//32] scales into the expected format:
https://github.com/drisspg/transformer_nuggets/blob/382cb0f19a5f615827174289b8ef552419d51fea/transformer_nuggets/mx/to_blocked.py#L11
This was surprisingly hard fought and would not have been possible w/ @albanD 😊

This allows this PR: #1625 to not have any dependencies on PT core updates while we add the required dtypes and bindings to cublas: pytorch/pytorch#145562

Follow up

Config needs more tuning

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 Infra failures during checkout of PyTorch

❌ 1 New Failure

As of commit 7473aca with merge base b2fb664 ():

NEW FAILURE - The following job has failed:

• Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_ao_drisspg/stack/31_3.9_cpu_

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

[vkuzo]

nice! if CI is green - looks good! I think this should have at least one numerical test though. Can be a follow-up PR if needed.

[albanD]

Very cool!