[Common][PyTorch] EP dispatch with unfused MXFP8 quantization by phu0ngng · Pull Request #3270 · NVIDIA/TransformerEngine

phu0ngng · 2026-07-28T20:33:22Z

Description

This PR adds MXFP8 support to the dispatch op of the NCCL EP path. The dispatch op is used in two places, and MXFP8 applies to both:

Dispatch forward bfloat16 tokens are quantized to MXFP8 internally and dispatched to the target experts; recv is returned as a per-expert GroupedTensor.
Combine backward the result-grad is scattered back to expert positions through the same (reverse) dispatch op, quantized to MXFP8, returning the expert-output grad as a per-expert GroupedTensor.

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

PyTorch frontend (`transformer_engine/pytorch/ep.py`, `distributed.py`, `csrc/extensions/ep.cpp`)**

The dispatch op quantizes bfloat16 tokens to MXFP8 internally when the buffer's dispatch_quant_recipe is set (MXFP8BlockScaling only for now); dispatch-forward recv is returned as a per-expert GroupedTensor. A pre-quantized input is rejected.
Combine backward reuses the dispatch op to scatter the result-grad: it quantizes the grad to MXFP8 and returns the expert-output grad as a per-expert GroupedTensor. Combine forward is unchanged (high-precision).
Recv data and block scales share a single caller-supplied (optionally symm-mem-backed) buffer, sliced into data-then-scale regions; the same convention is used for the combine backward grad buffer.

Common backend (`common/ep/ep_backend.cpp`, `include/.../ep.h`, `comm_window.h`)**

Backend and public headers extended to carry block-scale buffers/windows through the dispatch primitive.

NCCL EP submodule**

Bumped 3rdparty/nccl-extensions to the revision providing block-scaled dispatch.

Tests (`tests/cpp_distributed/test_ep.cu`, `tests/pytorch/distributed/run_ep.py`, `run_test_ep.sh`)**

Added C++ distributed coverage for the MXFP8 dispatch path.
Added PyTorch MXFP8 test passes for dispatch forward (normal, zero-copy, eager IO modes) and combine backward, gated behind a dedicated NVTE_EP_MXFP8_PASS run since the grouped path pins the per-expert alignment process-wide.

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

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

greptile-apps · 2026-07-28T20:36:00Z

Greptile Summary

Adds unfused MXFP8 quantization to NCCL expert-parallel dispatch and combine backward.

Routes MXFP8 data and block scales together through the common backend and PyTorch bindings.
Returns dispatched activations and combine-backward gradients as per-expert grouped tensors.
Adds symmetric-memory buffer slicing, pool teardown support, and distributed MXFP8 coverage.

Confidence Score: 4/5

The MXFP8 implementation appears safe to merge, with only the previously reported repository ignore-file deletion still requiring attention.

The EP data-and-scale paths are consistently represented across the Python frontend, binding layer, common backend, and distributed tests; however, the repository ignore configuration remains entirely deleted, leaving generated and local artifacts visible to Git.

Files Needing Attention: .gitignore

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ep.py	Adds recipe-controlled MXFP8 quantization, shared data/scale output allocation, grouped-tensor construction, and quantized combine backward.
transformer_engine/pytorch/csrc/extensions/ep.cpp	Extends PyTorch EP bindings to validate and forward compact MXFP8 scale-inverse tensors and symmetric-memory windows.
transformer_engine/common/ep/ep_backend.cpp	Builds NCCL descriptors for block scales and enables scale-forwarding dispatch for MXFP8 payloads.
transformer_engine/pytorch/distributed.py	Adds explicit teardown for the process-wide symmetric-memory pool and its cached NCCL-backed segments.
tests/pytorch/distributed/run_ep.py	Adds dedicated distributed coverage for MXFP8 dispatch, caller-provided buffers, zero-copy behavior, eager mode, and combine backward.

Sequence Diagram

sequenceDiagram
  participant User
  participant EP as PyTorch EP
  participant Quant as MXFP8 Quantizer
  participant Bind as C++ Binding
  participant Backend as NCCL EP Backend
  User->>EP: ep_dispatch(bfloat16 tokens)
  EP->>Quant: Quantize to E4M3 + E8M0 scales
  Quant-->>EP: Data and compact scale-inverse
  EP->>Bind: Dispatch data and scales
  Bind->>Backend: Block-scaled dispatch
  Backend-->>EP: Expert-major data and scales
  EP-->>User: Per-expert GroupedTensor
  User->>EP: ep_combine(expert outputs)
  EP-->>User: High-precision combined result
  User->>EP: Backward(result gradient)
  EP->>Quant: Quantize result gradient
  EP->>Backend: Reverse dispatch data and scales
  Backend-->>User: Per-expert grouped output gradient

_{Reviews (3): Last reviewed commit: "Separate EP dispatch-forward and combine..." | Re-trigger Greptile}

phu0ngng · 2026-07-28T23:38:03Z

        alignment: int = 0,
        payload_dtype: torch.dtype = torch.bfloat16,
        device: Optional[torch.device] = None,
+        dispatch_quant_recipe: Optional["Recipe"] = None,


Let me separate the recipe for dispatch fwd and combine bwd, in case users want to use different quantization for fwd and bwd in the future.

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng · 2026-07-30T17:42:07Z

/te-ci L1 pytorch

YangFei1990 · 2026-07-30T23:51:03Z

Is this intentional?

YangFei1990 · 2026-08-01T18:02:55Z

+        (MXFP8 for now), ``tokens`` is the quantized tensor kept as the autograd operand so grad
+        reaches the pre-quant input. Recv outputs are carved/allocated here: a caller may supply
+        ``recv_tokens`` / ``recv_topk_weights``, else they are sized to ``rows``."""
+        from .quantized_tensor import QuantizedTensor


nit: why not import at top?

YangFei1990 · 2026-08-01T18:16:14Z

        alignment: int = 0,
        payload_dtype: torch.dtype = torch.bfloat16,
        device: Optional[torch.device] = None,
+        dispatch_fwd_quant_recipe: Optional["Recipe"] = None,


Do we expect user/mcore to provide a recipe? I think recipe is specific concept in TE that we do not want to let end user to touch. Ideally, mcore can provide something dtype/quant config, e.g.

ep_dtype: if we allow single dtype for all

dispatch_dtype, combine_dtype: if we want to allow different dtypes for dispatch and combine (you know combine can be tricky)

dispatch_fwd_dtype, dispatch_bwd_dtype, combine_fwd_dtype, combine_bwd_dtype: fine-grained full config

I would recommend the middle option, as combine and dispatch might want to use different recipe, but between dispatch they might want to use same. The arge can take things like mxfp8/nvfp4, and TE will automatically pick the recipe

YangFei1990 · 2026-08-01T18:19:04Z

        # Host mirror of total_recv_tokens, set by ep_prepare in eager mode.
        self._host_total_recv_tokens: Optional[int] = None
+        # Set by ep_prepare; dispatch/combine read the routing it cached in handle_mem.
+        self.prepared = False


Seems this is only in the raw functions for assertion? What is the usage?

YangFei1990 · 2026-08-01T18:23:24Z

    return torch.empty(*shape, dtype=dtype, device=device)


+def _quantize(x: torch.Tensor, recipe):


I feel it would be better to find a better home for these quantization functions, seems to me they do not EP specific, or the requirement is implicit.

YangFei1990 · 2026-08-01T18:32:59Z

+            if tokens._fp8_dtype != tex.DType.kFloat8E4M3:
+                raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.")
+            # recv data + scales share one buffer (data then scales); carve or allocate it here.
+            recv_tokens, recv_scale_inv = _scale_alloc_io(


What if recv_tokens is not None in this case? Basically how do we handle user allocated symm buffer with MXFP8 dispatch? Mcore has persistent buffers allocated, would be good to have TE expose the scale allocate APIs for Mcore can call as well.

YangFei1990 · 2026-08-02T05:10:42Z

+    if buffer.dispatch_fwd_quant_recipe is not None:
+        # The grouped recv's packed scales only match the grouped layout when each expert slot is
+        # 128-aligned.
+        if buffer.alignment <= 0 or buffer.alignment % 128 != 0:


Is this a requirement for NCCL EP? The downstream group gemms have their own asserts

YangFei1990 · 2026-08-02T05:16:55Z

    return _SYMM_MEM_POOL


+def release_symm_mem_pool(device: Optional[torch.device] = None) -> None:


This is great, shall we also call it in ep_finalize?

YangFei1990 · 2026-08-02T05:35:05Z

+    # and their NCCL windows deregister. Fail loudly if this internal has changed shape,
+    # otherwise the windows would silently leak past destroy_process_group().
+    if torch_cached:
+        pools = getattr(symm_mem, "_symm_mem_pools", None)


Do we release all the pools, not only the one we created?

YangFei1990 · 2026-08-02T05:38:44Z

+        )
+        return data, scale_inv
+    data = _alloc_io((rows, data_cols), data_dtype, device, zero_copy)
+    scale_inv = _alloc_io((rows, scale_cols), scale_dtype, device, zero_copy)


Currently symm pool based allocation/deallocation will break cuda graph. We should either

Reuse the current persistent buffer design

Push PT to fix graph issue and migrate all symm allocation to pool based

EP Dispatch with MXFP8

f31b881

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng requested review from ksivaman and ptrendx as code owners July 28, 2026 20:33

restore gitignore

5c58e3b

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng requested review from YangFei1990 and vthumbe1503 July 28, 2026 20:35

phu0ngng requested a review from zhongbozhu July 28, 2026 23:33

phu0ngng commented Jul 28, 2026

View reviewed changes

phu0ngng added 2 commits July 30, 2026 10:29

EP Dispatch with MXFP8

c5ce029

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Separate EP dispatch-forward and combine-backward quant recipes

041b6e6

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

YangFei1990 requested changes Aug 2, 2026

View reviewed changes

		return torch.empty(*shape, dtype=dtype, device=device)


		def _quantize(x: torch.Tensor, recipe):

		return _SYMM_MEM_POOL


		def release_symm_mem_pool(device: Optional[torch.device] = None) -> None:

Uh oh!

Conversation

phu0ngng commented Jul 28, 2026

Description

Type of change

Changes

PyTorch frontend (transformer_engine/pytorch/ep.py, distributed.py, csrc/extensions/ep.cpp)**

Common backend (common/ep/ep_backend.cpp, include/.../ep.h, comm_window.h)**

NCCL EP submodule**

Tests (tests/cpp_distributed/test_ep.cu, tests/pytorch/distributed/run_ep.py, run_test_ep.sh)**

Checklist:

Uh oh!

greptile-apps Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phu0ngng commented Jul 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

PyTorch frontend (`transformer_engine/pytorch/ep.py`, `distributed.py`, `csrc/extensions/ep.cpp`)**

Common backend (`common/ep/ep_backend.cpp`, `include/.../ep.h`, `comm_window.h`)**

Tests (`tests/cpp_distributed/test_ep.cu`, `tests/pytorch/distributed/run_ep.py`, `run_test_ep.sh`)**

greptile-apps Bot commented Jul 28, 2026 •

edited

Loading