[PyTorch] Reduce CPU overhead on eager EP dispatch and combine

[phu0ngng]

Description

Lowers the per-call CPU cost of the eager NCCL EP dispatch and combine autograd paths. Numerics and public API are unchanged. Eager ops keep the CPU near the critical path (the dispatch forward must read the per-step recv count on the host before sizing recv), so Python and dispatch overhead was exposed rather than hidden behind async GPU work.

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

• Fuse eager prepare and dispatch into one C++ op and unify the eager and non-eager paths through it. The recv-count read (a host sync on pinned UVA memory, no D2H copy) and the recv-output allocation now happen inside the C++ op, so no Python runs between the count read and the dispatch launch.
• Shorter autograd operand lists: _EpPrepareAndDispatch and _EpCombine take only the differentiable tensors as direct operands. The non-differentiable buffer tensors ride on the EpBuffer object, and handle_mem is stashed as a plain ctx attribute instead of save_for_backward.
• Eager bypass of torch.library: eager is not graph-capturable, so combine forward, combine backward, and dispatch backward call the backend ops directly. Non-eager paths keep the transformer_engine_ep::* custom ops unchanged.
• Cache ep_get_zero_copy() on ctx instead of re-querying it in every combine backward.

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 reduces eager expert-parallel CPU overhead by fusing prepare and dispatch at the C++ boundary, shortening autograd operand lists, and bypassing torch.library for eager backend calls.

• Adds a fused native prepare-and-dispatch binding that dynamically allocates eager receive outputs.
• Unifies eager and fixed-capacity dispatch through _EpPrepareAndDispatch.
• Moves buffer-owned state onto EpBuffer and autograd context attributes.
• Adds eager MXFP8 dispatch forward/backward coverage.

Confidence Score: 4/5

The PR is not yet safe to merge because eager expert-parallel dispatch still breaks torch.compile tracing or full-graph compilation.

The eager path directly invokes tex.ep_prepare_and_dispatch, while only the non-eager path uses the registered custom operation required for Dynamo-compatible tracing, so the previously reported compilation failure remains.

Files Needing Attention: transformer_engine/pytorch/ep.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ep.py	Unifies dispatch autograd paths and directly invokes native eager operations, but eager dispatch remains incompatible with full-graph compilation.
transformer_engine/pytorch/csrc/extensions/ep.cpp	Adds the fused native prepare-and-dispatch implementation, eager receive-count synchronization, dynamic output allocation, and pybind registration.
transformer_engine/pytorch/csrc/extensions.h	Declares the fused prepare-and-dispatch extension interface and its optional receive buffers.
tests/pytorch/distributed/run_ep.py	Adds eager MXFP8 dispatch autograd coverage that checks gradients through dispatch backward and the quantizer STE.

Sequence Diagram

sequenceDiagram
  participant Caller
  participant EP as ep_dispatch
  participant Autograd as _EpPrepareAndDispatch
  participant Native as ep_prepare_and_dispatch
  participant Backend as NCCL EP backend
  Caller->>EP: tokens, routing, EpBuffer
  EP->>EP: optionally quantize MXFP8
  EP->>Autograd: apply(...)
  alt eager mode
    Autograd->>Native: direct pybind call
    Native->>Backend: prepare routing
    Native->>Native: synchronize and read recv count
    Native->>Native: allocate recv outputs
    Native->>Backend: dispatch
  else fixed-capacity mode
    Autograd->>Autograd: allocate/reuse static outputs
    Autograd->>Native: torch.library custom op
    Native->>Backend: prepare and dispatch
  end
  Autograd-->>Caller: recv tokens, weights, expert counts

Loading

_{Reviews (3): Last reviewed commit: "[PyTorch] Call the backend directly for ..." | Re-trigger Greptile}

[phu0ngng]

/te-ci L1 pytorch