diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6e09ed316f..567fb7d48b 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Multi-process PyTorch EP tests, launched via torchrun (one process per GPU).""" +from contextlib import nullcontext import os import sys import unittest @@ -11,6 +12,8 @@ import torch import torch.distributed as dist +import transformer_engine.pytorch as te +from transformer_engine.pytorch import ops as te_ops from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, @@ -25,6 +28,9 @@ _ep_combine_raw, _ep_dispatch_raw, ) +from transformer_engine.pytorch.ep_reference import BlockScaledTensor, MoeEpReference, MoeFormat +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp, _pack_grouped_linear_weights +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" @@ -39,6 +45,7 @@ # NVTE_EP_TOKENS_PER_RANK. HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512")) TOP_K = 2 +INTERMEDIATE_DIM = 16 TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32")) @@ -136,6 +143,31 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _make_moe_inputs(rank, ep_size, device="cuda"): + """Deterministic BF16 activations and FP32 top-k router weights.""" + generator = torch.Generator(device=device) + generator.manual_seed(2026 + rank) + tokens = ( + torch.randn( + TOKENS_PER_RANK, + HIDDEN_DIM, + generator=generator, + dtype=torch.float32, + device=device, + ) + * 0.25 + ).to(torch.bfloat16) + router_logits = torch.randn( + TOKENS_PER_RANK, + ep_size * NUM_LOCAL_EXPERTS, + generator=generator, + dtype=torch.float32, + device=device, + ) + topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) + return topk_idx, tokens, torch.softmax(topk_logits, dim=-1) + + def _degroup_mxfp8(recv_grouped, valid_counts=None): """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert @@ -146,6 +178,20 @@ def _degroup_mxfp8(recv_grouped, valid_counts=None): return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) +def _reference_weights(op): + """Pack GroupedLinear weights into MoeEpReference ``(E, in, out)`` layout.""" + packed = _pack_grouped_linear_weights(op, block_scaled_cls=BlockScaledTensor) + if isinstance(packed, torch.Tensor): + return packed.detach() + return BlockScaledTensor( + data=packed.data.detach(), + scale=packed.scale.detach(), + format=packed.format, + logical_shape=packed.logical_shape, + axis=packed.axis, + ) + + class _Cfg: rank: int world_size: int @@ -277,6 +323,66 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) + def _make_moe_model( + self, + *, + fuse_ops=True, + intermediate_dim=INTERMEDIATE_DIM, + recipe=None, + ): + """Build an EP MoE Sequential. + + With ``fuse_ops=True``, dispatch routing extras stay internal so + MegaMoE can claim the sequence when its gates pass. With + ``fuse_ops=False``, those extras are returned to the caller and the + Sequential stays unfused. Pass an MXFP8 ``recipe`` to construct + natively quantized GroupedLinear weights via ``quantized_model_init``. + """ + buffer = self._make_buffer() + dispatch = te_ops.Dispatch(buffer) + init_ctx = ( + te.quantized_model_init(enabled=True, recipe=recipe) + if recipe is not None + else nullcontext() + ) + previous_single_param = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" + try: + with init_ctx: + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + ) + finally: + if previous_single_param is None: + del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] + else: + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param + combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) + combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) + fc1.set_extra_input_channel(0, "tokens_per_expert") + activation.set_extra_input_channel(0, "routing_weights") + fc2.set_extra_input_channel(0, "tokens_per_expert") + return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 + # Prepare @_eager_test_include @@ -528,6 +634,154 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) + @_eager_test_include + def test_bf16_moe_sequential_vs_reference(self): + """Fused BF16 MegaMoE Sequential matches a BF16 ``MoeEpReference``.""" + self._run_moe_sequential_vs_reference(quantization=None) + + @_eager_test_include + def test_mxfp8_moe_sequential_vs_reference(self): + """Fused MegaMoE Sequential under MXFP8 autocast matches the QDQ reference.""" + self._run_moe_sequential_vs_reference(quantization="mxfp8") + + def _run_moe_sequential_vs_reference(self, *, quantization): + """Compare Sequential to ``MoeEpReference``. + + Sequential uses MegaMoE when supported. The MXFP8 reference QDQ GEMM + operands to MXFP8 and matmuls in FP32; the BF16 reference runs FP32 + GEMMs with no QDQ. + """ + if not EAGER: + self.skipTest("variable-size reference comparison requires eager EP mode") + if torch.cuda.get_device_capability() != (10, 7): + self.skipTest("MegaMoE fusion requires Rubin SM107") + try: + from cudnn.moe_ep import MoeEp # noqa: F401 + except ImportError as exc: + self.skipTest(f"cudnn.moe_ep.MoeEp is not installed ({type(exc).__name__}: {exc})") + + # MegaMoE SM107 requires intermediate_size % 256 == 0. + intermediate_dim = 256 + recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None + model, fc1, fc2 = self._make_moe_model( + fuse_ops=True, intermediate_dim=intermediate_dim, recipe=recipe + ) + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(3100 + self.cfg.rank) + with torch.no_grad(): + for op in (fc1, fc2): + weights = op.weight.quantized_tensors + if weights is None: + weights = op.weight.split_into_quantized_tensors() + for expert in range(NUM_LOCAL_EXPERTS): + weight = ( + torch.randn( + weights[expert].shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + weights[expert].copy_(weight) + if quantization == "mxfp8": + self.assertIsInstance(weights[expert], MXFP8Tensor) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + seq_tokens = tokens.detach().clone().requires_grad_(True) + seq_topk_weights = topk_weights.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() + ) + with autocast_ctx: + seq_out = model(seq_tokens, topk_idx, seq_topk_weights) + + forward_ops = model._module_groups[0]._forward_ops + self.assertEqual(len(forward_ops), 1) + self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + self.assertIsInstance(seq_out, torch.Tensor) + self.assertEqual(seq_out.dtype, torch.bfloat16) + + fc1_weight = _reference_weights(fc1) + fc2_weight = _reference_weights(fc2) + for op, packed in ((fc1, fc1_weight), (fc2, fc2_weight)): + packed_data = packed.data if isinstance(packed, BlockScaledTensor) else packed + self.assertEqual(packed_data.data_ptr(), op.weight.rowwise_data.data_ptr()) + self.assertTrue(packed_data.permute(0, 2, 1).is_contiguous()) + if isinstance(packed, BlockScaledTensor): + self.assertEqual(packed.scale.data_ptr(), op.weight.scale_inv.data_ptr()) + self.assertTrue(packed.scale.permute(0, 2, 1).is_contiguous()) + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=intermediate_dim, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + output_format=MoeFormat.BF16, + combine_format=MoeFormat.BF16, + apply_topk_in_fc1=True, + generate_c=True, + compute_dtype=torch.float32 if quantization == "mxfp8" else torch.bfloat16, + gemm_format=MoeFormat.MXFP8 if quantization == "mxfp8" else MoeFormat.BF16, + ) + ref_out, fc1_c, route_metadata = reference( + tokens.detach(), + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + ) + self.assertEqual(ref_out.dtype, torch.bfloat16) + + dy = ( + torch.randn( + seq_out.shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + seq_out.backward(dy) + grad_tokens, grad_fc1, grad_fc2, grad_topk_weights = reference.backward( + dy, + tokens.detach(), + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + fc1_c, + route_metadata, + ) + torch.cuda.synchronize() + if quantization == "bf16": + tolerances = {"rtol": 1.6e-2, "atol": 5e-5} + else: + tolerances = {"rtol": 0.125, "atol": 0.25} + torch.testing.assert_close(seq_out, ref_out, **tolerances) + torch.testing.assert_close( + seq_tokens.grad, grad_tokens.to(dtype=seq_tokens.dtype), **tolerances + ) + torch.testing.assert_close(seq_topk_weights.grad, grad_topk_weights.float(), **tolerances) + for op, ref_grad in ((fc1, grad_fc1), (fc2, grad_fc2)): + seq_grad = op.weight.grad + self.assertEqual(seq_grad.dtype, torch.bfloat16) + self.assertEqual( + tuple(seq_grad.shape), + (NUM_LOCAL_EXPERTS, op.out_features, op.in_features), + ) + self.assertTrue(seq_grad.is_contiguous()) + torch.testing.assert_close( + seq_grad, + ref_grad.transpose(1, 2).to(dtype=seq_grad.dtype), + **tolerances, + ) + @_zero_copy_test_include @_mxfp8_align_test def test_combine_bwd_mxfp8_caller_grad_out(self): diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 2799bfdf5d..08cac08ee1 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -26,6 +26,7 @@ __all__ = [ "EpBuffer", "ep_bootstrap", + "get_ep_group", "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", @@ -183,6 +184,11 @@ def is_ep_bootstrapped() -> bool: return _BOOTSTRAPPED +def get_ep_group() -> Optional[dist.ProcessGroup]: + """Return the process group registered by :func:`ep_bootstrap`.""" + return _EP_GROUP + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py new file mode 100644 index 0000000000..d38c6a80bf --- /dev/null +++ b/transformer_engine/pytorch/ep_reference.py @@ -0,0 +1,991 @@ +# Copyright (c) 2028 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. + +The implementation deliberately favors readable semantics over performance. It +supports a one-rank execution path and a variable-size ``all_to_all_single`` EP +path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. + +The quantized tensor layouts are logical (unswizzled) layouts. A production +kernel may reorder scale factors internally without changing this API contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +class MoeFormat(str, Enum): + """Public and communication formats supported by the reference.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = value + return tuple(result) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Portable data-plus-scale representation for MXFP8 or NVFP4. + + ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` + has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with + two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` + replaces that axis by one scale per block. + """ + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + fmt = _parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + shape = tuple(int(dim) for dim in self.logical_shape) + if not shape or any(dim < 0 for dim in shape): + raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") + axis = _normalize_axis(self.axis, len(shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", shape) + object.__setattr__(self, "axis", axis) + self._validate_storage() + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + def _validate_storage(self) -> None: + if self.data.device != self.scale.device: + raise ValueError("block-scaled data and scale must be on the same device") + + logical_extent = self.logical_shape[self.axis] + scale_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, self.block_size), + ) + if tuple(self.scale.shape) != scale_shape: + raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") + + if self.format is MoeFormat.MXFP8: + expected_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + data_shape = self.logical_shape + if self.data.dtype != expected_dtype: + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) + else: + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: + raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") + data_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, 2), + ) + + if tuple(self.data.shape) != data_shape: + raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") + if self.scale.dtype != expected_scale_dtype: + raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return the logical tensor with block scales applied.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + packed = packed.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" + + levels = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=values.device, + ) + magnitudes = values.abs().unsqueeze(-1) + distances = (magnitudes - levels).abs() + minimum = distances.amin(dim=-1, keepdim=True) + candidates = distances == minimum + codes = torch.arange(8, dtype=torch.int64, device=values.device) + any_code = torch.where(candidates, codes, 8).amin(dim=-1) + even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) + magnitude_code = torch.where(even_code < 8, even_code, any_code) + sign_code = torch.signbit(values).to(torch.int64) << 3 + return magnitude_code | sign_code + + +def quantize_blockwise( + tensor: torch.Tensor, + format: Union[MoeFormat, str], + *, + axis: int = -1, +) -> BlockScaledTensor: + """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. + + MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward + positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and + E4M3 scales rounded to nearest. + """ + + fmt = _parse_format(format) + if fmt is MoeFormat.BF16: + raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") + if not tensor.is_floating_point(): + raise TypeError(f"tensor must be floating point, got {tensor.dtype}") + + axis = _normalize_axis(axis, tensor.ndim) + logical_shape = tuple(tensor.shape) + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + block_count = _ceil_div(logical_extent, block_size) + padded_extent = block_count * block_size + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) + + value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 + scale_float = blocks.abs().amax(dim=-1) / value_limit + if fmt is MoeFormat.MXFP8: + safe_scale = torch.where(scale_float > 0, scale_float, 1.0) + scale_float = torch.where( + scale_float > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(scale_float), + ) + scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + scale_dtype = _require_torch_dtype("float8_e4m3fn") + + scale = scale_float.to(scale_dtype) + scale_for_math = scale.float() + reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) + + if fmt is MoeFormat.MXFP8: + data_dtype = _require_torch_dtype("float8_e4m3fn") + data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] + else: + codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) + low = codes[..., 0::2] + high = codes[..., 1::2] + data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] + + return BlockScaledTensor( + data=data.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format=fmt, + logical_shape=logical_shape, + axis=axis, + ) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +@dataclass(frozen=True) +class _DispatchPlan: + """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" + + send_expert: torch.Tensor # local expert id per sent route + send_weight: torch.Tensor # router weight per sent route + send_token_idx: torch.Tensor # source token per sent route + send_slot_idx: torch.Tensor # source top-k slot per sent route + send_counts: Tuple[int, ...] # routes sent to each rank + recv_counts: Tuple[int, ...] # routes received from each rank + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +def _decode_tensor( + tensor: MoeTensor, + *, + name: str, + expected_shape: Tuple[int, ...], + quantized_axis: int, + dtype: torch.dtype, +) -> torch.Tensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.logical_shape != expected_shape: + raise ValueError( + f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" + ) + if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): + raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") + return tensor.dequantize(dtype=dtype) + + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") + if not tensor.is_floating_point(): + raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") + return tensor.to(dtype=dtype) + + +def _format_round_trip( + tensor: torch.Tensor, + format: MoeFormat, + *, + dtype: torch.dtype, +) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).to(dtype) + return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) + + +def _qdq( + tensor: torch.Tensor, + format: Optional[MoeFormat], + *, + axis: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Simulate an MXFP8/NVFP4 GEMM operand in PyTorch. + + The MegaMoE kernel quantizes once and feeds the quantized values into the + GEMM (FP32 accumulate). PyTorch has no MXFP8 matmul, so the reference + round-trips through block-scale storage and dequantizes back to ``dtype``. + """ + if format is None or format is MoeFormat.BF16: + return tensor.to(dtype=dtype) + return quantize_blockwise(tensor, format, axis=axis).dequantize(dtype=dtype) + + +def _swiglu_scale_fp32( + gate: torch.Tensor, + up: torch.Tensor, + weights: torch.Tensor, + *, + apply_scale: bool, + dtype: torch.dtype, +) -> torch.Tensor: + """SiLU(gate)*up[*weights] in fp32, then one cast to ``dtype``. + + Matches ``tex.scaled_swiglu``: promote, multiply, and store once. Stepwise + BF16 ``F.silu(g) * up * w`` extra-rounds between those muls. + """ + out = F.silu(gate.float()) * up.float() + if apply_scale: + out = out * weights.float() + return out.to(dtype=dtype) + + +class MoeEpReference: + """Reference implementation of routed SwiGLU experts plus EP dispatch. + + Global experts are assigned contiguously: rank ``r`` owns + ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit + initialized process group for multi-rank execution; ``None`` means a + one-rank reference even if the default distributed group is initialized. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + compute_dtype: torch.dtype = torch.float32, + gemm_format: Optional[Union[MoeFormat, str]] = None, + ) -> None: + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError( + f"compute_dtype must be torch.float32 or torch.bfloat16, got {compute_dtype}" + ) + + if ep_group is None: + ep_size, ep_rank = 1, 0 + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "ep_group requires an initialized torch.distributed process group" + ) + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size != 0: + raise ValueError( + f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" + ) + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + self.compute_dtype = compute_dtype + self.gemm_format = None if gemm_format is None else _parse_format(gemm_format) + if self.gemm_format is MoeFormat.BF16: + self.gemm_format = None + + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = ( + 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + ) + if hidden_size % required_multiple != 0: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" + f" {name}={fmt.value}" + ) + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(" + f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " + f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " + f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " + f"output={self.output_format.value}, combine={self.combine_format.value}, " + f"gemm_format={None if self.gemm_format is None else self.gemm_format.value}, " + f"compute_dtype={self.compute_dtype})" + ) + + def _gemm_dtype(self) -> torch.dtype: + """FP32 accumulate when simulating an MXFP8 GEMM; otherwise ``compute_dtype``.""" + return torch.float32 if self.gemm_format is MoeFormat.MXFP8 else self.compute_dtype + + def _stage_gemm_operand( + self, + tensor: torch.Tensor, + *, + already_quantized: bool, + axis: int, + ) -> torch.Tensor: + """Apply the kernel's pre-GEMM quantize, then dequant for a PyTorch matmul.""" + dtype = self._gemm_dtype() + if already_quantized or self.gemm_format is None: + return tensor.to(dtype=dtype) + return _qdq(tensor, self.gemm_format, axis=axis, dtype=dtype) + + def _collective_device(self, device: torch.device) -> torch.device: + """Device the process group can run ``all_to_all_single`` on. + + Gloo only implements all-to-all for CPU tensors, so CUDA tensors are + staged through host memory; NCCL groups communicate in place. + """ + if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return send_counts.clone() + comm_device = self._collective_device(send_counts.device) + staged = send_counts.to(comm_device) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single(recv_counts, staged, group=self.ep_group) + return recv_counts.to(send_counts.device) + + def _all_to_all( + self, + send: torch.Tensor, + send_counts: Sequence[int], + recv_counts: Sequence[int], + ) -> torch.Tensor: + if self.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + output_shape = (sum(recv_counts), *send.shape[1:]) + recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) + dist.all_to_all_single( + recv, + staged, + output_split_sizes=list(recv_counts), + input_split_sizes=list(send_counts), + group=self.ep_group, + ) + return recv.to(send.device) + + def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: + """Route valid ``topk_idx`` entries to destination ranks, stably by rank. + + Backward reuses this so gradient re-dispatch reproduces the exact + forward route order. + """ + + device = topk_idx.device + token_count = topk_idx.shape[0] + flat_expert = topk_idx.reshape(-1).to(torch.int64) + flat_weight = topk_weights.reshape(-1).float() + valid = flat_expert != -1 + invalid_negative = flat_expert < -1 + invalid_high = flat_expert >= self.num_experts + if bool((invalid_negative | invalid_high).any().item()): + bad = flat_expert[invalid_negative | invalid_high][0].item() + raise ValueError(f"topk_idx contains out-of-range expert id {bad}") + + flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) + flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) + expert = flat_expert[valid] + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + order = torch.argsort(destination, stable=True) + + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), + send_weight=flat_weight[valid].index_select(0, order), + send_token_idx=flat_token[valid].index_select(0, order), + send_slot_idx=flat_slot[valid].index_select(0, order), + send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), + recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), + ) + + def _run_local_experts( + self, + tokens: torch.Tensor, + local_expert_idx: torch.Tensor, + route_weight: torch.Tensor, + fc1_weight: torch.Tensor, + fc2_weight: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output = torch.empty( + (tokens.shape[0], self.hidden_size), + dtype=self.compute_dtype, + device=tokens.device, + ) + fc1_c_rows = [] if self.generate_c else None + for expert in range(self.experts_per_rank): + positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + expert_tokens = tokens.index_select(0, positions) + gate_up = expert_tokens @ fc1_weight[expert] + if fc1_c_rows is not None: + # Raw pre-SwiGLU accumulator: before clamp, no router weight. + fc1_c_rows.append(gate_up.to(torch.bfloat16)) + gate, up = gate_up.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + gate = gate.float().clamp(max=self.gate_up_clamp) + up = up.float().clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + weights = route_weight.index_select(0, positions).unsqueeze(-1) + # FP32 SiLU*up[*scale], then one cast. MegaMoE then quantizes this + # intermediate for the FC2 MXFP8 GEMM; the reference QDQ mimics that. + intermediate = _swiglu_scale_fp32( + gate, + up, + weights, + apply_scale=self.apply_topk_in_fc1, + dtype=self._gemm_dtype(), + ) + if self.gemm_format is not None: + intermediate = _qdq( + intermediate, + self.gemm_format, + axis=-1, + dtype=self._gemm_dtype(), + ) + expert_output = intermediate @ fc2_weight[expert] + if not self.apply_topk_in_fc1: + expert_output = (expert_output.float() * weights.float()).to( + dtype=self.compute_dtype + ) + expert_output = _format_round_trip( + expert_output, + self.combine_format, + dtype=self.compute_dtype, + ) + output.index_copy_(0, positions, expert_output) + fc1_c = None + if fc1_c_rows is not None: + fc1_c = ( + torch.cat(fc1_c_rows) + if fc1_c_rows + else torch.empty( + (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device + ) + ) + return output, fc1_c + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: + """Run dispatch, local experts, return routing, top-k reduce, and encode. + + Shapes: + activation: ``(T, H)`` + fc1_weight: ``(E_local, H, 2 * I)`` + fc2_weight: ``(E_local, I, H)`` + topk_idx/topk_weights: ``(T, K)`` + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when constructed with ``generate_c=True``. ``fc1_c`` is the BF16 + pre-SwiGLU FC1 accumulator of every route this rank's experts + processed, ``(local_routes, 2 * I)``, grouped by local expert and + ordered within each expert by (source rank, source token-major route + order); captured before the gate/up clamp, without the router weight. + ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies + the route behind ``fc1_c`` row ``i`` for the backward gradient + re-dispatch. + """ + + if topk_idx.ndim != 2: + raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") + token_count = topk_idx.shape[0] + route_shape = (token_count, self.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) + + device = _tensor_device(activation) + inputs = { + "fc1_weight": _tensor_device(fc1_weight), + "fc2_weight": _tensor_device(fc2_weight), + "topk_idx": topk_idx.device, + "topk_weights": topk_weights.device, + } + for name, input_device in inputs.items(): + if input_device != device: + raise ValueError(f"{name} must be on {device}, got {input_device}") + + activation_float = self._stage_gemm_operand( + _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(activation, BlockScaledTensor), + axis=-1, + ) + fc1_float = self._stage_gemm_operand( + _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=( + self.experts_per_rank, + self.hidden_size, + 2 * self.intermediate_size, + ), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc1_weight, BlockScaledTensor), + axis=1, + ) + fc2_float = self._stage_gemm_operand( + _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc2_weight, BlockScaledTensor), + axis=1, + ) + + plan = self._dispatch_plan(topk_idx, topk_weights) + send_token_idx = plan.send_token_idx + send_slot_idx = plan.send_slot_idx + send_counts, recv_counts = plan.send_counts, plan.recv_counts + send_tokens = activation_float.index_select(0, send_token_idx) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( + dtype=self.compute_dtype + ) + + route_metadata = None + if self.generate_c: + recv_src_rank = torch.repeat_interleave( + torch.arange(self.ep_size, device=device), + torch.tensor(recv_counts, device=device), + ) + recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) + recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) + # Stable sort by local expert reproduces the fc1_c row order + # (grouped by expert; source order preserved within each group). + fc1_c_order = torch.argsort(recv_expert, stable=True) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) + # recv rows are ordered by source rank, then that source's token-major + # route order, so the per-expert position grouping below realizes the + # documented fc1_c ordering. + recv_output, fc1_c = self._run_local_experts( + recv_tokens, + recv_expert, + recv_weight, + fc1_float, + fc2_float, + ) + + returned = self._all_to_all(recv_output, recv_counts, send_counts) + # Combine payload is combine_format (BF16 round-trip above); reduce in + # fp32 so top-k summation does not extra-round in compute_dtype. + combine_plane = torch.zeros( + (token_count * self.top_k, self.hidden_size), + dtype=torch.float32, + device=device, + ) + send_flat_slot = send_token_idx * self.top_k + send_slot_idx + combine_plane.index_copy_(0, send_flat_slot, returned.float()) + reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) + + if self.output_format is MoeFormat.BF16: + output = reduced.to(torch.bfloat16) + else: + output = quantize_blockwise(reduced, self.output_format, axis=-1) + if self.generate_c: + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward pass consuming the ``generate_c=True`` stash. + + ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, + and the FC2 input are all rebuilt from it, so no post-SwiGLU forward + intermediate needs to be saved. ``route_metadata`` alone reconstructs + the mapping between re-dispatched rows and ``fc1_c`` rows and drives + the gradient return scatter. + + Quantization round-trips (input decode, ``combine_format``, + ``output_format``) are treated as straight-through identities; + ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. + + Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, + grad_topk_weights)``. Activation and expert weight gradients use + ``compute_dtype``. SwiGLU and router-weight scaling recompute in + fp32 and round once into ``compute_dtype`` before the FC2 GEMMs, + matching ``tex.scaled_swiglu``. The returned ``grad_topk_weights`` + remain float32. + """ + + if not self.generate_c: + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) + token_count = topk_idx.shape[0] + if tuple(grad_output.shape) != (token_count, self.hidden_size): + raise ValueError( + f"grad_output shape must be {(token_count, self.hidden_size)}, got" + f" {tuple(grad_output.shape)}" + ) + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + + device = _tensor_device(activation) + two_i = 2 * self.intermediate_size + activation_float = self._stage_gemm_operand( + _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(activation, BlockScaledTensor), + axis=-1, + ) + fc1_float = self._stage_gemm_operand( + _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc1_weight, BlockScaledTensor), + axis=1, + ) + fc2_float = self._stage_gemm_operand( + _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ), + already_quantized=isinstance(fc2_weight, BlockScaledTensor), + axis=1, + ) + if fc1_c.shape != (int(route_metadata.shape[0]), two_i): + raise ValueError( + f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" + f" {tuple(fc1_c.shape)}" + ) + + # Re-dispatch the FC1 inputs, router weights, and output gradients + # along the identical forward routes. + plan = self._dispatch_plan(topk_idx, topk_weights) + send_counts, recv_counts = plan.send_counts, plan.recv_counts + grad_output_float = grad_output.to(dtype=self.compute_dtype) + recv_tokens = self._all_to_all( + activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + # Same FP32-on-wire / compute_dtype-for-activation cast as forward. + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( + dtype=self.compute_dtype + ) + recv_grad = self._all_to_all( + grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + + # route_metadata rows are in fc1_c order; sorting them by + # (src_rank, src_token, src_slot) reproduces the receive order, giving + # the permutation between re-dispatched rows and fc1_c rows. + metadata = route_metadata.to(device=device, dtype=torch.int64) + local_routes = metadata.shape[0] + if local_routes > 0: + token_span = int(metadata[:, 2].max().item()) + 1 + recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] + perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j + else: + perm = torch.empty((0,), dtype=torch.int64, device=device) + x_rows = torch.empty_like(recv_tokens) + x_rows.index_copy_(0, perm, recv_tokens) + w_rows = torch.empty_like(recv_weight) + w_rows.index_copy_(0, perm, recv_weight) + dy_rows = torch.empty_like(recv_grad) + dy_rows.index_copy_(0, perm, recv_grad) + + c_rows = fc1_c.to(dtype=self.compute_dtype) + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=self.compute_dtype, + device=device, + ) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + grad_fc1 = torch.zeros_like(fc1_float) + grad_fc2 = torch.zeros_like(fc2_float) + for expert in range(self.experts_per_rank): + positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + c = c_rows.index_select(0, positions) + x = x_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1).float() + d_y = dy_rows.index_select(0, positions) + + gate, up = c.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + g = gate.float().clamp(max=self.gate_up_clamp) + u = up.float().clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate.float(), up.float() + sig = torch.sigmoid(g) + s = g * sig + h = s * u + + if self.apply_topk_in_fc1: + h_fc2 = (h * w).to(dtype=self._gemm_dtype()) + d_y_pre = d_y.to(dtype=self._gemm_dtype()) + else: + h_fc2 = h.to(dtype=self._gemm_dtype()) + d_y_pre = (d_y.float() * w).to(dtype=self._gemm_dtype()) + if self.gemm_format is not None: + h_fc2 = _qdq(h_fc2, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) + d_y_pre = _qdq(d_y_pre, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) + grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre + d_h_fc2 = (d_y_pre @ fc2_float[expert].transpose(0, 1)).float() + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = (d_y.float() * (h @ fc2_float[expert].float())).sum(dim=-1) + + d_g = d_h * u * (sig * (1 + g * (1 - sig))) + d_u = d_h * s + if self.gate_up_clamp is not None: + d_gate = d_g * (gate.float() <= self.gate_up_clamp) + up_f = up.float() + d_up = d_u * ((up_f >= -self.gate_up_clamp) & (up_f <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1).to(dtype=self._gemm_dtype()) + if self.gemm_format is not None: + d_c = _qdq(d_c, self.gemm_format, axis=-1, dtype=self._gemm_dtype()) + grad_fc1[expert] = x.transpose(0, 1) @ d_c + d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) + + # Return the route gradients to their source ranks and scatter-add. + returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + grad_activation = torch.zeros( + (token_count, self.hidden_size), + dtype=torch.float32, + device=device, + ) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx.float()) + grad_activation = grad_activation.to(dtype=self.compute_dtype) + grad_topk_weights = torch.zeros( + (token_count * self.top_k,), dtype=torch.float32, device=device + ) + grad_topk_weights.index_copy_( + 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw + ) + return ( + grad_activation, + grad_fc1, + grad_fc2, + grad_topk_weights.view(token_count, self.top_k), + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "quantize_blockwise", +] diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index caab9c9c8d..29ebc1bf4b 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -22,7 +22,9 @@ from .all_reduce import AllReduce from .basic_linear import BasicLinear from .bias import Bias +from .combine import Combine from .constant_scale import ConstantScale +from .dispatch import Dispatch from .dropout import Dropout from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported from .identity import Identity diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py new file mode 100644 index 0000000000..4d4740c517 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel combine operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, is_symm_backed +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_grad_buffer( + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError("grad_out must be contiguous.") + if tensor.requires_grad: + raise ValueError("grad_out must not require gradients.") + return tensor + + +class Combine(BasicOperation): + """Combine pre-weighted local expert outputs with NCCL EP. + + The operation uses routing state produced by a :class:`Dispatch` with the + same :class:`EpBuffer`. + """ + + def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: + super().__init__() + self.buffer = buffer + self.num_local_tokens = ( + buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) + ) + if self.num_local_tokens < 0: + raise ValueError("num_local_tokens must be non-negative.") + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Combine input must have shape (R, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + + expert_out = input_ + if self.buffer.zero_copy: + expert_out = _alloc_io( + tuple(input_.shape), + input_.dtype, + input_.device, + True, + ) + expert_out.copy_(input_) + + result = torch.empty( + self.num_local_tokens, + self.buffer.hidden_dim, + dtype=input_.dtype, + device=input_.device, + ) + torch.ops.transformer_engine_ep.combine( + self.buffer.handle_mem, + expert_out, + result, + ) + + if ctx.requires_grad: + grad_out = kwargs.get("grad_out") + if self.buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes combine gradients per step and cannot use " + "a caller-supplied grad_out" + ) + grad_out = _validate_grad_buffer( + grad_out, + shape=tuple(input_.shape), + dtype=input_.dtype, + device=input_.device, + ) + if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): + raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") + ctx.grad_out = grad_out + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.save_for_backward(self.buffer.handle_mem) + + return result + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + grad_output, + grad_input, + ) + return grad_input, () diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py new file mode 100644 index 0000000000..9762ed76f7 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel dispatch operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, ep_prepare +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_output_buffer( + name: str, + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous.") + if tensor.requires_grad: + raise ValueError(f"{name} must not require gradients.") + return tensor + + +class Dispatch(BasicOperation): + """Dispatch BF16 tokens to local experts with NCCL EP. + + The extra inputs are routing indices and FP32 routing weights. The extra + outputs are local tokens-per-expert and received routing weights. + """ + + num_extra_inputs: int = 2 + num_extra_outputs: int = 2 + + def __init__(self, buffer: EpBuffer) -> None: + super().__init__() + self.buffer = buffer + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}.") + expected_route_shape = (input_.shape[0], self.buffer.top_k) + if tuple(topk_idx.shape) != expected_route_shape: + raise ValueError( + f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." + ) + if tuple(topk_weights.shape) != expected_route_shape: + raise ValueError( + f"topk_weights shape must be {expected_route_shape}, " + f"got {tuple(topk_weights.shape)}." + ) + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): + if tensor.device != input_.device: + raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") + + recv_tokens = kwargs.get("recv_tokens") + recv_topk_weights = kwargs.get("recv_topk_weights") + if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): + raise ValueError( + "eager mode sizes dispatch outputs per step and cannot use " + "caller-supplied receive buffers" + ) + + tokens_per_expert = ep_prepare(self.buffer, topk_idx) + rows = ( + self.buffer._host_total_recv_tokens + if self.buffer.eager + else self.buffer.recv_capacity_per_rank + ) + if rows is None: + raise RuntimeError("NCCL EP dispatch receive size is unavailable.") + rows = int(rows) + recv_shape = (rows, self.buffer.hidden_dim) + recv_tokens = _validate_output_buffer( + "recv_tokens", + recv_tokens, + shape=recv_shape, + dtype=self.buffer.payload_dtype, + device=self.buffer.device, + ) + recv_topk_weights = _validate_output_buffer( + "recv_topk_weights", + recv_topk_weights, + shape=(rows,), + dtype=torch.float32, + device=self.buffer.device, + ) + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (rows,), + torch.float32, + self.buffer.device, + self.buffer.zero_copy, + ) + + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.save_for_backward(self.buffer.handle_mem) + + return recv_tokens, [(tokens_per_expert, recv_topk_weights)] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + + grad_recv_weights = basic_op_grad_extra_outputs[0][1] + if grad_recv_weights is None: + grad_recv_weights = torch.zeros( + grad_output.shape[0], + dtype=torch.float32, + device=grad_output.device, + ) + else: + grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous() + + grad_input = torch.empty( + ctx.input_shape, + dtype=ctx.input_dtype, + device=grad_output.device, + ) + grad_topk_weights = torch.empty( + ctx.topk_weights_shape, + dtype=torch.float32, + device=grad_output.device, + ) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + grad_output, + grad_recv_weights, + grad_input, + grad_topk_weights, + ) + return grad_input, [()], [(None, grad_topk_weights)] diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index dc9dcd6dc3..7aa45a89f5 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -35,3 +35,4 @@ GroupedMLP_CuTeGEMMGLU, GroupedMLP_CuTeGEMMUnary, ) +from .moe_ep import FusedMoeEp # pylint: disable=wrong-import-position diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py new file mode 100644 index 0000000000..14747a726a --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -0,0 +1,439 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MegaMoE-backed expert-parallel MoE fusion (cudnn.moe_ep.MoeEp).""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any, Optional + +import torch + +from ...constants import MXFP8_BLOCK_SCALING_SIZE +from ...ep import get_ep_group +from ...quantization import Recipe +from ...tensor import GroupedTensor, Quantizer +from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU +from ..fuser import register_forward_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext + + +def _grouped_weight(op: GroupedLinear) -> GroupedTensor: + """Return the single packed ``(E, out, in)`` parameter required by MegaMoE.""" + if not op.single_grouped_weight or not isinstance(op.weight, GroupedTensor): + raise ValueError("FusedMoeEp requires GroupedLinear(single_grouped_weight=True)") + return op.weight + + +def _pack_grouped_linear_weights(op: GroupedLinear, *, block_scaled_cls: Optional[type] = None): + """View a packed TE ``(E, out, in)`` weight as MegaMoE ``(E, in, out)``. + + The permutation is intentionally not made contiguous. MegaMoE internally + permutes back to ``(E, out, in)`` before requesting contiguous storage, so + retaining this view lets that request reuse GroupedLinear's original packed + buffer. ``block_scaled_cls`` selects the ``BlockScaledTensor`` type (cuDNN + MegaMoE vs the PyTorch reference). + """ + weight = _grouped_weight(op) + num_groups = op.num_groups + out_features = op.out_features + in_features = op.in_features + expected_data_numel = num_groups * out_features * in_features + if weight.rowwise_data is None or weight.rowwise_data.numel() != expected_data_numel: + raise ValueError( + "GroupedLinear weight must have compact rowwise storage with shape " + f"({num_groups}, {out_features}, {in_features})" + ) + + data_nk = weight.rowwise_data.view(num_groups, out_features, in_features) + if weight.quantizer is not None: + recipe = weight.quantizer._get_compatible_recipe() + if recipe is None or not recipe.mxfp8(): + raise TypeError("FusedMoeEp only supports dense BF16 or MXFP8 grouped weights") + if weight._with_gemm_swizzled_scales: + raise NotImplementedError( + "FusedMoeEp requires unswizzled MXFP8 weight scales, got a GEMM-swizzled tensor" + ) + if in_features % MXFP8_BLOCK_SCALING_SIZE != 0: + raise ValueError( + f"MXFP8 weight K={in_features} is not divisible by {MXFP8_BLOCK_SCALING_SIZE}" + ) + scale_cols = in_features // MXFP8_BLOCK_SCALING_SIZE + expected_scale_numel = num_groups * out_features * scale_cols + if weight.scale_inv is None or weight.scale_inv.numel() != expected_scale_numel: + raise ValueError( + "GroupedLinear MXFP8 scales must have compact rowwise storage with shape " + f"({num_groups}, {out_features}, {scale_cols})" + ) + if block_scaled_cls is None: + from cudnn.moe_ep import BlockScaledTensor as block_scaled_cls + packed_data = data_nk.view(torch.float8_e4m3fn).permute(0, 2, 1) + packed_scale = ( + weight.scale_inv.view(num_groups, out_features, scale_cols) + .view(torch.float8_e8m0fnu) + .permute(0, 2, 1) + ) + return block_scaled_cls( + data=packed_data, + scale=packed_scale, + format="mxfp8", + logical_shape=tuple(packed_data.shape), + axis=1, + ) + return data_nk.permute(0, 2, 1) + + +def _flatten_moe_weight(weight) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Split a MegaMoE weight into tensors that ``save_for_backward`` can hold.""" + if isinstance(weight, torch.Tensor): + return weight, None + return weight.data, weight.scale + + +def _restore_moe_weight(data: torch.Tensor, scale: Optional[torch.Tensor], block_scaled_cls: type): + """Rebuild the MegaMoE weight saved by :func:`_flatten_moe_weight`.""" + if scale is None: + return data + return block_scaled_cls( + data=data, + scale=scale, + format="mxfp8", + logical_shape=tuple(data.shape), + axis=1, + ) + + +def _grouped_weight_grad(op: GroupedLinear, grad: torch.Tensor) -> list[Optional[torch.Tensor]]: + """Convert a MegaMoE ``(E, in, out)`` wgrad to one TE ``(E, out, in)`` grad.""" + weight = _grouped_weight(op) + expected_shape = (op.num_groups, op.in_features, op.out_features) + if tuple(grad.shape) != expected_shape: + raise RuntimeError( + f"MegaMoE weight gradient must have shape {expected_shape}, got {tuple(grad.shape)}" + ) + if not weight.requires_grad: + return [None] + + # copy_ performs the float32-to-parameter-dtype conversion while writing + # directly into the contiguous layout expected by the grouped parameter. + param_grad = torch.empty( + (op.num_groups, op.out_features, op.in_features), + dtype=weight.dtype, + device=grad.device, + ) + param_grad.copy_(grad.transpose(1, 2)) + return [param_grad] + + +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weight = op.weight if op.single_grouped_weight else None + weight_ok = ( + isinstance(weight, GroupedTensor) + and weight.dtype is torch.bfloat16 + and weight.rowwise_data is not None + ) + if weight_ok and weight.quantizer is not None: + recipe = weight.quantizer._get_compatible_recipe() + weight_ok = ( + recipe is not None + and recipe.mxfp8() + and not weight._with_gemm_swizzled_scales + and weight.rowwise_data is not None + and weight.scale_inv is not None + ) + + return ( + not op.use_bias + and not op._scale_bias + and op.single_grouped_weight + and not op.single_grouped_bias + and not op._accumulate_into_main_grad + and not op._is_distributed_weight() + and not op.wgrad_store.delay_wgrad_compute() + and weight_ok + ) + + +def _import_cudnn_moe_ep(): + """Return ``cudnn.moe_ep.MoeEp`` or ``None`` if the package is missing.""" + try: + from cudnn.moe_ep import MoeEp + except ImportError: + return None + return MoeEp + + +def _routing_extras_internal( + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`cudnn.moe_ep.MoeEp`, so it can only replace the + sequence when those two outputs feed exactly these ops and are not + returned to the caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + +def _megamoe_supported(buffer, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: + """Static MegaMoE capability gates that can be checked before first launch.""" + if _import_cudnn_moe_ep() is None: + return False + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + return False + if buffer.max_tokens_per_rank is None or buffer.max_tokens_per_rank <= 0: + return False + if buffer.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: + return False + if buffer.top_k > 32: + return False + ep_group = get_ep_group() + ep_size = 1 if ep_group is None else ep_group.size() + if ep_size > 16: + return False + return True + + +def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: + if len(window) != 5: + return False + if recipe is not None and not recipe.mxfp8(): + return False + dispatch, fc1, activation, fc2, combine = window + if not ( + isinstance(dispatch, Dispatch) + and isinstance(fc1, GroupedLinear) + and isinstance(activation, ScaledSwiGLU) + and isinstance(fc2, GroupedLinear) + and isinstance(combine, Combine) + ): + return False + if dispatch.buffer is not combine.buffer: + return False + buffer = dispatch.buffer + if not buffer.eager or buffer.payload_dtype is not torch.bfloat16: + return False + if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): + return False + if activation.activation_recompute_in_mlp or activation.glu_interleave_size is not None: + return False + if not _routing_extras_internal(dispatch, fc1, activation, fc2): + return False + if not _megamoe_supported(buffer, fc1, fc2): + return False + return ( + fc1.num_groups == buffer.num_local_experts + and fc2.num_groups == buffer.num_local_experts + and fc1.in_features == buffer.hidden_dim + and fc2.out_features == buffer.hidden_dim + and fc1.out_features == 2 * fc2.in_features + ) + + +class FusedMoeEp(FusedOperation): + """Joint EP MoE fusion implemented with :class:`cudnn.moe_ep.MoeEp`.""" + + def __init__( + self, + *, + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, + combine: Combine, + ) -> None: + super().__init__([dispatch, fc1, activation, fc2, combine]) + moe_ep_cls = _import_cudnn_moe_ep() + if moe_ep_cls is None: + raise ImportError( + "FusedMoeEp requires cudnn.moe_ep.MoeEp. Install the in-tree " + "cuDNN frontend with: pip install --force-reinstall " + "'./cudnn_frontend[moe_ep]'" + ) + from cudnn.moe_ep import BlockScaledTensor + + ep_group = get_ep_group() + ep_size = 1 if ep_group is None else ep_group.size() + self._block_scaled_cls = BlockScaledTensor + self._moe = moe_ep_cls( + num_experts=dispatch.buffer.num_local_experts * ep_size, + hidden_size=dispatch.buffer.hidden_dim, + intermediate_size=fc2.in_features, + top_k=dispatch.buffer.top_k, + ep_group=ep_group, + max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, + apply_topk_in_fc1=True, + generate_c=True, + ) + + @property + def dispatch(self) -> Dispatch: + return self.basic_ops[0] + + @property + def fc1(self) -> GroupedLinear: + return self.basic_ops[1] + + @property + def fc2(self) -> GroupedLinear: + return self.basic_ops[3] + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") + if any(kwargs for kwargs in basic_op_kwargs): + raise NotImplementedError("FusedMoeEp does not support per-operation output buffers.") + + topk_idx, topk_weights = basic_op_extra_inputs[0] + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + with torch.no_grad(): + fc1_weight = _pack_grouped_linear_weights( + self.fc1, block_scaled_cls=self._block_scaled_cls + ) + fc2_weight = _pack_grouped_linear_weights( + self.fc2, block_scaled_cls=self._block_scaled_cls + ) + output, fc1_c, route_metadata = self._moe( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + + if any(ctx.requires_grad for ctx in basic_op_ctxs): + fc1_data, fc1_scale = _flatten_moe_weight(fc1_weight) + fc2_data, fc2_scale = _flatten_moe_weight(fc2_weight) + basic_op_ctxs[0].save_for_backward( + input_, + fc1_data, + fc1_scale, + fc2_data, + fc2_scale, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + # Dispatch extras are channel-bound with output_to_caller=False and are + # only consumed by ops inside this fusion, so they need not be materialized. + return output, [ + (None, None), + (), + (), + (), + (), + ] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + ( + input_, + fc1_data, + fc1_scale, + fc2_data, + fc2_scale, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) = basic_op_ctxs[0].saved_tensors + fc1_weight = _restore_moe_weight(fc1_data, fc1_scale, self._block_scaled_cls) + fc2_weight = _restore_moe_weight(fc2_data, fc2_scale, self._block_scaled_cls) + grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._moe.backward( + grad_output, + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + # MegaMoE returns float32 grads in (E, in, out); each GroupedLinear has + # one packed (E, out, in) parameter. + fc1_param_grads = _grouped_weight_grad(self.fc1, grad_fc1) + fc2_param_grads = _grouped_weight_grad(self.fc2, grad_fc2) + return ( + grad_input.to(dtype=input_.dtype), + [(), fc1_param_grads, (), fc2_param_grads, ()], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), ()], + ) + + +def fuse_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused: Any, +) -> list[FusibleOperation]: + """Fuse supported five-op EP MoE sequences into MegaMoE. + + Unfused Sequential (NCCL dispatch/combine + GroupedLinear + ScaledSwiGLU) + is the default. MegaMoE is claimed only when ``_matches`` succeeds. + """ + del unused + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + window = ops[idx : idx + 5] + if _matches(window, recipe): + dispatch, fc1, activation, fc2, combine = window + out.append( + FusedMoeEp( + dispatch=dispatch, + fc1=fc1, + activation=activation, + fc2=fc2, + combine=combine, + ) + ) + idx += 5 + else: + out.append(ops[idx]) + idx += 1 + return out + + +register_forward_backward_fusion(fuse_ops, prepend=True) + + +__all__ = ["FusedMoeEp"]