[Pytorch][Do not merge] Let MXFP8 dispatch return plain torch tensor

[YangFei1990]

Description

This PR allows TE's EP dispatch API return plain torch tensor when doing MXFP8 dispatch. Previously it will return a grouped tensor

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR changes MXFP8 expert-parallel dispatch and combine-backward outputs from grouped quantized tensors to opaque plain PyTorch tensor carriers.

• Packs E4M3 data and E8M0 scales into payload-shaped carrier storage.
• Adds a helper that reconstructs the per-expert grouped MXFP8 view.
• Updates distributed EP tests and fused grouped-MLP carrier handling.

Confidence Score: 4/5

The PR does not appear safe to merge because fused grouped-MLP carrier reconstruction remains disabled by attributes that are never assigned.

MXFP8 dispatch and combine backward now return packed plain tensors, while the fused grouped-MLP reconstruction branches still default their unassigned carrier flags to false; the packed E4M3 data and scales can therefore enter computation as ordinary BF16 values.

Files Needing Attention: transformer_engine/pytorch/ops/fused/grouped_mlp.py, transformer_engine/pytorch/ep.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ep.py	Replaces grouped MXFP8 dispatch and combine gradients with opaque plain-tensor carriers and adds grouped-view reconstruction.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds carrier reconstruction branches, but the existing unset-flag issue still leaves those branches unreachable.
tests/pytorch/distributed/run_ep.py	Updates MXFP8 EP tests to validate plain carrier identity, storage aliasing, reconstruction, and numerical results.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[BF16 tokens] --> B[MXFP8 quantization]
  B --> C[EP dispatch]
  C --> D[Plain tensor carrier<br/>E4M3 data + E8M0 scales]
  D --> E[mxfp8_carrier_to_grouped]
  E --> F[Grouped MXFP8 tensor]
  F --> G[Fused grouped MLP]

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into mxfp8_dispatch_..." | Re-trigger Greptile}

[greptile-apps]

Carrier helper omitted from exports

The ep_dispatch documentation directs callers to mxfp8_carrier_to_grouped, but the helper is absent from the module's __all__ list. Wildcard imports therefore omit the function required to consume the new opaque representation, forcing callers to rely on an undeclared implementation detail.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!