[NCCL EP] Add MXFP8 dispatch support

megatron/core/fp8_utils.py

@@ -347,11 +347,9 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, F
     try:
         custom_recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=quantizer_factory)
     except AttributeError:
-        raise ValueError(
-            """CustomRecipe recipe is not available in this version of 
+        raise ValueError("""CustomRecipe recipe is not available in this version of 
             Transformer Engine. Please make sure you are using TE version 
-            >= 2.9.0.dev0."""
-        )
+            >= 2.9.0.dev0.""")
     return custom_recipe
 
 
@@ -811,6 +809,26 @@ def get_fp8_recipe(config: TransformerConfig):
             )
         return fp8_recipe
 
+    def get_fp8_recipe_for_a2a(a2a_dtype: str):
+        """Return the fp8 recipe for quantizing an MoE dispatch/combine (a2a) payload over
+        the wire, or None for a high-precision wire.
+
+        Dedicated helper rather than get_fp8_recipe: the wire payload is E4M3
+        activations/activation-grads in both directions, independent of the compute recipe
+        (whose format selection and compute-only knobs like fp8_dpa do not apply to a
+        communication payload).
+
+        Arguments:
+            a2a_dtype (str): Wire dtype, 'bf16' (returns None) or 'mxfp8'.
+        """
+        if a2a_dtype == 'bf16':
+            return None
+        if a2a_dtype == 'mxfp8':
+            return transformer_engine.common.recipe.MXFP8BlockScaling(
+                fp8_format=transformer_engine.common.recipe.Format.E4M3
+            )
+        raise ValueError(f"Unsupported a2a wire dtype: {a2a_dtype!r}.")
+
     def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
         """Return fp8 context manager.
 
@@ -900,6 +918,14 @@ def get_fp8_recipe(config: TransformerConfig):
         """Returns None since TE is not available."""
         return None
 
+    def get_fp8_recipe_for_a2a(a2a_dtype: str):
+        """Raises for a quantized wire dtype since TE is not available."""
+        if a2a_dtype == 'bf16':
+            return None
+        raise RuntimeError(
+            f"a2a wire dtype {a2a_dtype!r} requires TransformerEngine, which is not available."
+        )
+
     def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
         """Returns dummy fp8 context manager since TE is not available."""
         return nullcontext()

megatron/core/transformer/moe/experts.py

@@ -533,6 +533,13 @@ def register_grouped_linear_params(
         register_grouped_linear_params(
             op, self.linear_fc1, fc1_single_grouped_weight, fc1_single_grouped_bias
         )
+        # FP8 dispatch: the FC1 input arrives as an opaque MXFP8 carrier tensor (TE EP dispatch
+        # packs E4M3 data + scales into a plain tensor's storage); tell TE to rebuild the
+        # grouped view at the op boundary.
+        # Only works with TE PR https://github.com/NVIDIA/TransformerEngine/pull/3355
+        # TODO: remove after TE support the grouped tensor path
+        if getattr(self.config, 'moe_dispatch_fwd_dtype', 'bf16') == 'mxfp8':
+            op.ep_mxfp8_carrier_input = True
         ops.append(op)
 
         # Activation and post-multiply probs (SwiGLU, clamped GLU, or SReLU).
@@ -634,6 +641,13 @@ def register_grouped_linear_params(
         register_grouped_linear_params(
             op, self.linear_fc2, fc2_single_grouped_weight, fc2_single_grouped_bias
         )
+        # FP8 combine backward: the FC2 output grad arrives as an opaque MXFP8 carrier tensor
+        # (TE EP combine backward, same packing as dispatch); tell TE to rebuild the grouped
+        # view at the op boundary.
+        # Only works with TE PR https://github.com/NVIDIA/TransformerEngine/pull/3355
+        # TODO: remove after TE support the grouped tensor path
+        if getattr(self.config, 'moe_combine_bwd_dtype', 'bf16') == 'mxfp8':
+            op.ep_mxfp8_carrier_grad = True
         ops.append(op)
 
         # Emulate submodule pre-forward hooks

megatron/core/transformer/moe/fused_a2a.py

@@ -3,6 +3,7 @@
 # Copyright (c) 2025 DeepSeek
 # Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
 
+import inspect
 from typing import Optional
 
 from megatron.core.utils import internal_api
@@ -692,27 +693,51 @@ def alloc_ep_symm_buffer(shape, dtype, ep_group):
         the per-call recv buffers it recycles."""
         return te_ep.symm_mem_alloc(shape, dtype, ep_group)
 
+    # Whether this TE build supports quantized EP dispatch/combine payloads.
+    _TE_EP_BUFFER_HAS_QUANT_RECIPE = (
+        "dispatch_fwd_quant_recipe" in inspect.signature(te_ep.EpBuffer).parameters
+    )
+
     def new_nccl_ep_buffer(
         top_k,
         max_tokens_per_rank,
         recv_capacity_per_rank,
         hidden_dim,
         num_local_experts,
         alignment=0,
+        dispatch_fwd_quant_recipe=None,
+        combine_bwd_quant_recipe=None,
     ):
         """Build a fresh TE EpBuffer for one dispatch/combine pair.
 
         The buffer owns handle_mem (the routing table dispatch writes and combine reads); a new one
         is built per dispatch and dropped after combine. Payload symm buffers are not owned here —
         they are caller-supplied to dispatch/combine or allocated on the fly by TE.
+
+        dispatch_fwd_quant_recipe / combine_bwd_quant_recipe (TE Recipe, optional): quantize the
+        dispatch forward payload / combine backward gradient over the wire,
+        the corresponding output is returned as a per-expert GroupedTensor.
         """
+        quant_kwargs = {}
+        if dispatch_fwd_quant_recipe is not None or combine_bwd_quant_recipe is not None:
+            if not _TE_EP_BUFFER_HAS_QUANT_RECIPE:
+                raise RuntimeError(
+                    "Quantized NCCL EP dispatch/combine payloads require a TransformerEngine "
+                    "build with EpBuffer quant-recipe support "
+                    "(dispatch_fwd_quant_recipe / combine_bwd_quant_recipe)."
+                )
+            quant_kwargs = dict(
+                dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe,
+                combine_bwd_quant_recipe=combine_bwd_quant_recipe,
+            )
         return te_ep.EpBuffer(
             top_k=top_k,
             max_tokens_per_rank=max_tokens_per_rank,
             recv_capacity_per_rank=recv_capacity_per_rank,
             hidden_dim=hidden_dim,
             num_local_experts=num_local_experts,
             alignment=alignment,
+            **quant_kwargs,
         )
 
     def nccl_ep_dispatch(

megatron/core/transformer/moe/token_dispatcher.py

@@ -10,6 +10,7 @@
 
 from megatron.core import utils
 from megatron.core.config import is_experimental_enabled
+from megatron.core.fp8_utils import get_fp8_recipe_for_a2a
 from megatron.core.fusions.fused_indices_converter import fused_indices_to_multihot
 from megatron.core.fusions.fused_pad_routing_map import fused_pad_routing_map
 from megatron.core.jit import jit_fuser
@@ -1587,6 +1588,9 @@ def __init__(
                         "ragged per-expert counts on device)."
                     )
 
+        self.dispatch_fwd_quant_recipe = get_fp8_recipe_for_a2a(config.moe_dispatch_fwd_dtype)
+        self.combine_bwd_quant_recipe = get_fp8_recipe_for_a2a(config.moe_combine_bwd_dtype)
+
         # Fresh EpBuffer per dispatch, held until the matching combine consumes it. dispatch
         # and combine share one buffer: handle_mem is the routing table that dispatch writes
         # and combine reads. Safe because dispatch i / combine i strictly alternate.
@@ -1705,6 +1709,8 @@ def dispatch(
             hidden_dim=self.hidden_dim,
             num_local_experts=self.num_local_experts,
             alignment=self.alignment,
+            dispatch_fwd_quant_recipe=self.dispatch_fwd_quant_recipe,
+            combine_bwd_quant_recipe=self.combine_bwd_quant_recipe,
         )
         # TE requires int64 indices and float32 weights.
         # token_indices/token_probs: [num_local_tokens, router_topk]

megatron/core/transformer/transformer_config.py

@@ -987,6 +987,20 @@ class TransformerConfig(ModelParallelConfig):
     symm-mem buffers are a fixed [recv_capacity, hidden] and cannot be resized per step) and the
     fused op (use_transformer_engine_op_fuser). Defaults to False."""
 
+    moe_dispatch_fwd_dtype: Literal['bf16', 'mxfp8'] = 'bf16'
+    """Wire dtype of the MoE dispatch forward payload ('ncclep' flex dispatcher only). With
+    'mxfp8', TransformerEngine quantizes the payload before the all-to-all and the receive
+    buffer comes back as a per-expert MXFP8 GroupedTensor that the grouped GEMM consumes
+    directly. Requires moe_grouped_gemm and use_transformer_engine_op_fuser. Defaults to
+    'bf16' (no quantization on the wire)."""
+
+    moe_combine_bwd_dtype: Literal['bf16', 'mxfp8'] = 'bf16'
+    """Wire dtype of the MoE combine backward gradient ('ncclep' flex dispatcher only). With
+    'mxfp8', TransformerEngine quantizes the gradient before the all-to-all and the
+    expert-output gradient comes back as a per-expert MXFP8 GroupedTensor that the grouped
+    GEMM backward consumes directly. Same requirements as moe_dispatch_fwd_dtype. Defaults to
+    'bf16' (no quantization on the wire)."""
+
     moe_mlp_glu_interleave_size: Optional[int] = None
     """When set, GLU activations in the MoE grouped MLP layer will use a
     block interleaved format. Instead of interpreting the input tensor
@@ -1737,6 +1751,22 @@ def __post_init__(self):
                     "or select the alltoall, DeepEP, or HybridEP dispatcher."
                 )
 
+        if self.moe_dispatch_fwd_dtype != 'bf16' or self.moe_combine_bwd_dtype != 'bf16':
+            if (
+                self.moe_token_dispatcher_type != "flex"
+                or self.moe_flex_dispatcher_backend != "ncclep"
+            ):
+                raise ValueError(
+                    "moe_dispatch_fwd_dtype / moe_combine_bwd_dtype require the 'ncclep' flex "
+                    "dispatcher backend."
+                )
+            if not (self.use_transformer_engine_op_fuser and self.moe_grouped_gemm):
+                raise ValueError(
+                    "moe_dispatch_fwd_dtype / moe_combine_bwd_dtype = 'mxfp8' require BOTH "
+                    "use_transformer_engine_op_fuser and moe_grouped_gemm: only the fused "
+                    "grouped GEMM path consumes the pre-quantized MXFP8 GroupedTensor payload."
+                )
+
         # moe_deepep_num_sms / moe_hybridep_num_sms are deprecated and unified into
         # moe_flex_dispatcher_num_sms. If either is set, route it (an explicit
         # moe_flex_dispatcher_num_sms takes precedence) and warn.

tests/unit_tests/a2a_overlap/test_schedule_layer_1f1b.py

@@ -27,6 +27,7 @@
     reset_model,
 )
 from tests.unit_tests.test_utilities import Utils
+from tests.unit_tests.transformer.moe.test_token_dispatcher import is_nccl_ep_fp8_dispatch_available
 
 # Transformer Engine 2.17 aborts in the A2A overlap suite with a pybind11 GIL dec_ref failure.
 pytestmark = pytest.mark.flaky_in_dev
@@ -476,15 +477,25 @@ def test_transformer_layer_overlap(self, dispatcher_type, flex_backend, fp8_flag
     @pytest.mark.skipif(
         not is_op_fuser_available(), reason="op-fuser (static-shape/zero-copy) needs TE>=2.14"
     )
-    def test_transformer_layer_overlap_zero_copy(self):
+    @pytest.mark.parametrize(
+        "wire_dtype", ["bf16", pytest.param("mxfp8", marks=pytest.mark.launch_on_gb200)]
+    )
+    def test_transformer_layer_overlap_zero_copy(self, wire_dtype):
         """ncclEP zero-copy under 1F1B a2a overlap must match the non-overlap reference.
 
         Zero-copy stays enabled in both runs, so this isolates the overlap schedule. It also
         compares the two ways zero-copy makes the dispatch-backward gradient symm-mem-backed:
         the reference gets it from the op-fuser's ``grad_input_buffer``, the overlap run from
         ``StageDispatchBwdGrad`` staging into the same buffer (plus the free_input symm guard).
-        bf16 op-fuser (SwiGLU, tp=1) -- no fp8/Blackwell dependency.
+        bf16 op-fuser (SwiGLU, tp=1) -- wire_dtype="bf16" has no fp8/Blackwell dependency.
+
+        wire_dtype="mxfp8" additionally sends the dispatch-fwd / combine-bwd payloads as the
+        opaque MXFP8 carrier. Both captures use the same wire dtype, so quantization is
+        common-mode; what this asserts is that the overlap schedule's staging/detach/free of
+        node-boundary tensors preserves the carrier bytes.
         """
+        if wire_dtype == "mxfp8" and not is_nccl_ep_fp8_dispatch_available():
+            pytest.skip("NCCL EP MXFP8 wire needs EpBuffer quant-recipe support and MXFP8 hardware")
         extra_kwargs = {}
         apply_flex_backend_kwargs(extra_kwargs, "flex", "ncclep")
         extra_kwargs.update(
@@ -496,6 +507,8 @@ def test_transformer_layer_overlap_zero_copy(self):
             gated_linear_unit=True,
             activation_func=F.silu,
             overlap_moe_expert_parallel_comm=True,
+            moe_dispatch_fwd_dtype=wire_dtype,
+            moe_combine_bwd_dtype=wire_dtype,
         )
         config = get_test_config(extra_kwargs=extra_kwargs)
         microbatches = 4

tests/unit_tests/models/test_hybrid_moe_model.py

@@ -187,6 +187,8 @@
     "moe_deepep_num_sms": None,
     "moe_enable_deepep": False,
     "moe_expert_capacity_factor": None,
+    "moe_combine_bwd_dtype": "bf16",
+    "moe_dispatch_fwd_dtype": "bf16",
     "moe_expert_rank_capacity_factor": None,
     "moe_ffn_hidden_size": 1856,
     "moe_flex_dispatcher_backend": "deepep",

tests/unit_tests/transformer/moe/test_paged_stashing.py

@@ -20,6 +20,7 @@
 from megatron.core.utils import is_te_min_version
 from megatron.training.initialize import _set_random_seed
 from tests.unit_tests.test_utilities import Utils
+from tests.unit_tests.transformer.moe.test_token_dispatcher import is_nccl_ep_fp8_dispatch_available
 
 # These tests configure mxfp8 + the TE op fuser, so they only run on Blackwell (sm100+). Mark the
 # whole module for the GB200 CI bucket (selection there is marker-driven; see
@@ -118,6 +119,8 @@ def __init__(
             moe_permute_fusion=kwargs.get("moe_permute_fusion", False),
             moe_flex_dispatcher_backend=kwargs.get("moe_flex_dispatcher_backend", None),
             moe_ncclep_zero_copy=kwargs.get("moe_ncclep_zero_copy", False),
+            moe_dispatch_fwd_dtype=kwargs.get("moe_dispatch_fwd_dtype", 'bf16'),
+            moe_combine_bwd_dtype=kwargs.get("moe_combine_bwd_dtype", 'bf16'),
             moe_grouped_gemm=kwargs.get("moe_grouped_gemm", False),
             moe_paged_stash=kwargs.get("moe_paged_stash", False),
             moe_expert_rank_capacity_factor=kwargs.get("moe_expert_rank_capacity_factor", None),
@@ -490,17 +493,24 @@ def teardown_method(self, method):
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     @pytest.mark.flaky_in_dev
     @pytest.mark.internal
-    def test_over_budget(self):
+    @pytest.mark.parametrize("wire_dtype", ["bf16", "mxfp8"])
+    def test_over_budget(self, wire_dtype):
         """Budget matches _NCCLEPManager._ensure_bootstrap; over_budget matches map-derived load.
 
         Mirrors TestPagedStashingOverBudget for HybridEP, plus the peak capacity NCCL EP
         reports -- HybridEP has no equivalent because it recovers by going dropless and so
         never needs to know how much was required. The capacity factor is deliberately below
         1.0: each rank receives num_tokens*topk on average, so 1.0 sits exactly at the mean
         and anything under it overflows.
+
+        wire_dtype="mxfp8" runs the same overflow accounting with MXFP8 dispatch-fwd /
+        combine-bwd wire payloads (the opaque carrier); the budget arithmetic under test is
+        payload-dtype independent, so the assertions are unchanged.
         """
         if not is_nccl_ep_available():
             pytest.skip("NCCL EP is not available")
+        if wire_dtype == "mxfp8" and not is_nccl_ep_fp8_dispatch_available():
+            pytest.skip("NCCL EP MXFP8 wire needs EpBuffer quant-recipe support and MXFP8 hardware")
 
         config.ENABLE_EXPERIMENTAL = True
 
@@ -526,6 +536,8 @@ def test_over_budget(self):
             moe_router_padding_for_quantization=True,
             gated_linear_unit=True,
             activation_func=F.silu,
+            moe_dispatch_fwd_dtype=wire_dtype,
+            moe_combine_bwd_dtype=wire_dtype,
         )
 
         seq_length = 1024
@@ -718,14 +730,20 @@ def run():
     @pytest.mark.flaky_in_dev
     @pytest.mark.internal
     @pytest.mark.parametrize("zero_copy", [False, True])
-    def test_forward_backward_4_layers(self, zero_copy):
+    @pytest.mark.parametrize("wire_dtype", ["bf16", "mxfp8"])
+    def test_forward_backward_4_layers(self, zero_copy, wire_dtype):
         """Test paged stashing with 4 MoE layers on ncclep static shape: two passes match.
 
-        zero_copy=True additionally exercises the ncclEP symm-mem zero-copy IO under paged stash."""
+        zero_copy=True additionally exercises the ncclEP symm-mem zero-copy IO under paged stash.
+        wire_dtype="mxfp8" additionally sends the dispatch-fwd / combine-bwd payloads as the
+        MXFP8 carrier; both passes use the same wire dtype, so quantization is common-mode and
+        the two-pass determinism tolerance is unchanged."""
         if not is_nccl_ep_available():
             pytest.skip("NCCL EP is not available")
         if zero_copy and not is_nccl_ep_zero_copy_available():
             pytest.skip("NCCL EP zero-copy TE API is not available")
+        if wire_dtype == "mxfp8" and not is_nccl_ep_fp8_dispatch_available():
+            pytest.skip("NCCL EP MXFP8 wire needs EpBuffer quant-recipe support and MXFP8 hardware")
 
         config.ENABLE_EXPERIMENTAL = True
 
@@ -752,6 +770,8 @@ def test_forward_backward_4_layers(self, zero_copy):
             gated_linear_unit=True,
             activation_func=F.silu,
             moe_ncclep_zero_copy=zero_copy,
+            moe_dispatch_fwd_dtype=wire_dtype,
+            moe_combine_bwd_dtype=wire_dtype,
         )
 
         seq_length = 1024