[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

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

Greptile Summary

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

Routes MXFP8 payload data and block scales through the common backend and PyTorch bindings.
Returns quantized dispatch and combine-backward outputs as per-expert grouped tensors.
Extends symmetric-memory buffer handling and teardown for zero-copy communication.
Adds dedicated distributed MXFP8 coverage for standard, eager, and zero-copy modes.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ep.py	Adds recipe-gated MXFP8 quantization, shared data/scale buffer slicing, grouped outputs, and quantized combine backward.
transformer_engine/pytorch/csrc/extensions/ep.cpp	Extends EP bindings to validate and forward MXFP8 data, scale tensors, and symmetric-memory window offsets.
transformer_engine/common/ep/ep_backend.cpp	Adds NCCL descriptors and dispatch configuration for routing MXFP8 block scales alongside token data.
transformer_engine/pytorch/distributed.py	Adds explicit symmetric-memory pool release bookkeeping for NCCL window teardown.
tests/pytorch/distributed/run_ep.py	Adds distributed dispatch-forward and combine-backward MXFP8 coverage across supported execution modes.
tests/cpp_distributed/test_ep.cu	Verifies that MXFP8 payload and scale rows follow the same expert-routing permutation.

Sequence Diagram

sequenceDiagram
  participant User
  participant PyEP as PyTorch EP
  participant Quant as MXFP8 Quantizer
  participant Bind as C++ Binding
  participant Backend as NCCL EP Backend
  User->>PyEP: ep_dispatch(BF16 tokens)
  PyEP->>Quant: quantize tokens
  Quant-->>PyEP: E4M3 data + E8M0 scales
  PyEP->>Bind: dispatch(data, scales)
  Bind->>Backend: route payload and scale rows
  Backend-->>PyEP: expert-major data + scales
  PyEP-->>User: GroupedTensor
  User->>PyEP: ep_combine backward
  PyEP->>Quant: quantize result gradient
  PyEP->>Backend: reverse dispatch(data, scales)
  Backend-->>User: grouped expert-output gradient

_{Reviews (4): Last reviewed commit: "bump nccl ep" | Re-trigger Greptile}

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

/te-ci L1 pytorch

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

Is this intentional?

Nope. Let me remove this change.

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: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?

For assertion only, but actually not very useful, so I will remove it.

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.

Removed the helper and inlined the impl.

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

This is not a requirement from NCCL EP.
You are right that it's the job of the GEMM ops to enforce their alignment requirements.

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?

release_symm_mem_pool is only required before destroy_process_group().
User may call ep_finalize without destroy_process_group().

Hmm why is that? I did not have it in mcore but it seems working (could be luck)

Also just see input device is never used, shall we remove?

On why

release_symm_mem_pool is only required before destroy_process_group().

I think destroy_process_group() deletes the NCCL communicator that is used to create the symmem, so the cached symmem-s become limbo.

Are you calling destroy_process_group() in MCore? Why?

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

…der CUDA graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

…CUDA-graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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

YangFei1990 · 2026-08-05T00:25:38Z

+            torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out)
+        else:
+            if tex.ep_get_zero_copy():
+                raise NotImplementedError(


What is the limitation here?

Stale. Let me remove it.

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

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

Comment thread transformer_engine/pytorch/ep.py Outdated

YangFei1990 requested changes Aug 2, 2026

View reviewed changes

phu0ngng added 9 commits August 4, 2026 08:52

EP Dispatch with MXFP8

56ba9ca

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

restore gitignore

b50d785

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

EP Dispatch with MXFP8

4ea01cf

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

Separate EP dispatch-forward and combine-backward quant recipes

d0a99f4

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

[PyTorch] Simplify EP quantize path and guard zero-copy pool alloc un…

75fb84d

…der CUDA graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[PyTorch] Add eager-mode coverage for MXFP8 EP combine backward

bfe103e

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

[PyTorch] Supply persistent symm-mem recv buffers for zero-copy 1F1B …

19cad25

…CUDA-graph capture Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

shorten quant enum

6a31417

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

bump nccl ep

778c2c8

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

phu0ngng force-pushed the ep_mxfp8 branch from 041b6e6 to 778c2c8 Compare August 4, 2026 16:29

YangFei1990 reviewed Aug 5, 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: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

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!

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!

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!

Uh oh!

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