[BugFix] Support DP/EP in AG/RS for FLASHINFER_CUTLASS FP8

[amirkl94]

Purpose

Add support for fp8 per-tensor DP and EP in MoE flashinfer cutlass path.
Changes:

1. Support fp8 hidden state communication in pynccl_wrapper .
2. Split the DP only case and EP case in prepare and finalize allgather object.

Test Plan

TBD

Test Result

TBD

[gemini-code-assist]

Code Review

This pull request introduces support for FP8 per-tensor DP and EP in the MoE FlashInfer Cutlass path. The changes correctly integrate the ncclFloat8e4m3 data type and update the from_torch method in pynccl_wrapper.py. Additionally, a new use_ep flag has been added to FlashInferAllGatherMoEPrepareAndFinalize and its factory function, create_flashinfer_prepare_finalize, to explicitly handle expert parallelism. The conditional logic in prepare and finalize methods has been updated to incorporate this new flag, ensuring proper differentiation between DP-only and DP+EP cases. The use_ep parameter is correctly propagated from the moe configuration, making the implementation robust. Overall, the changes align well with the stated purpose of supporting FP8 and explicit DP/EP handling.

[robertgshaw2-redhat]

FYI - @amirkl94 I also made this work in this PR by casting to int8, but your solution is cleaner

• [MoE Refactor] Integrate Naive Prepare Finalize into MK #32567

Could you make your change on top of #32567?

[amirkl94]

FYI - @amirkl94 I also made this work in this PR by casting to int8, but your solution is cleaner

• [MoE Refactor] Integrate Naive Prepare Finalize into MK #32567

Could you make your change on top of #32567?

Yes, I'm testing this now on top of your branch as well.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @amirkl94.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork