[mxfp8 moe training] add triton kernel for mxfp8 dequantization

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately #3203
• [mxfp8 moe training] improve unit tests #3201
• [mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all #3197
• ->[mxfp8 moe training] add triton kernel for mxfp8 dequantization #3195

---

[mxfp8 moe training] add triton kernel for mxfp8 dequantization

Summary

• Traces show dequantization kernel in mxfp8 a2a is slow, this PR adds a triton kernel for this which is much faster for large "M" (local_batch_size * seq_len) which is what we care about for MoE training.

Test plan

• pytest test/prototype/mx_formats/test_kernels.py -k mxfp8_dequant

Benchmarks

input_shape        torch_us    triton_us    torch_gbps    triton_gbps  triton_speedup
---------------  ----------  -----------  ------------  -------------  ----------------
(1, 8192, 7168)      36.864       39.968       4828.44        4453.46  0.922x
(2, 8192, 7168)     287.712       78.88        1237.32        4513.08  3.647x
(4, 8192, 7168)     560.32       150.56        1270.67        4728.9   3.722x
(8, 8192, 7168)    1110.9        297.984       1281.82        4778.67  3.728x

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 82ded0b with merge base b644211 ():

NEW FAILURE - The following job has failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  RuntimeError: Command docker exec -t e8c73dcaefd4f7e904527d0a3dfad256673515d454a741752ea47d38f5b4b2d9 /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/dtypes/test_nf4.py::TestComm::test_comm

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

[vkuzo]

lgtm, didn't look at the rest too closely

[drisspg]

Looks like there is some ptx for going to bf16

__CUDA_HOSTDEVICE_FP8_DECL__
__nv_bfloat16_raw __nv_cvt_e8m0_to_bf16raw(const __nv_fp8_storage_t x)
{
    __nv_bfloat16_raw res;

#if (__CUDA_FP8_INTERNAL_CAN_RELY_ON_PTX_FOR_SHORTTYPESCVT__)
    unsigned short in = (unsigned short)x;
    unsigned hr = 0U;
    asm("{cvt.rn.bf16x2.ue8m0x2 %0, %1;}\n"
                : "=r"(hr)
                : "h"(in));

    res.x = (unsigned short)hr;
#else
    res.x = __internal_e8m0_to_bf16(x);
#endif

    return res;
}

[danielvegamyhre]

Looks like there is some ptx for going to bf16

Sweet where is this from? I actually tried looking in TE for PTX examples for this but all I could find was casting fp32 -> e8m0 (for computing scale): https://github.com/NVIDIA/TransformerEngine/blob/dd9433e7ad28c12f27da9770be54c9c584e85fa0/transformer_engine/common/util/ptx.cuh#L134

Will try it out later