perf(MoE): Use TE quant/dequant for SwiGLU fp8 input store to improve performance and stability

[xiaoxi-wangfj]

Description

Replace native .to(fp8) casting in SwiGLU with Transformer Engine quant/dequant interfaces for storing activation inputs in FP8.

Benefits:

1. Higher performance – For quant+dequant operator performance, the dynamic quantization and dequantization method in transformer_engineprovides a 1.48x to 1.66x speedup compared to the native method.
2. Better numerical stability – The dynamic quantization and dequantization method in transformer_engine handles extreme values more gracefully. Native .to(fp8) can cause underflow of very small numbers to 0, or overflow of large values to inf, while TE scaling reduces these issues.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ksivaman]

LGTM, except a small nit to not import internal API as the exact structure of the files might change or move around in TE

[xiaoxi-wangfj]

@ksivaman Thanks for the review! I’ve addressed all the comments and updated the PR. Could you please re-review when you get a chance?

[Phlip79]

We are changing our review process and marking all open, unlabeled PRs as draft. This change will go in effect starting once #3659 is merged.

Moving forward, all PRs will be required to start as draft PRs. If you wish to get your PR merged, mark your PR as “Ready for review”. Read more about the new process at submit.md.