From c159fb64b92b05e79911edf20400396187402906 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 13 Jul 2026 21:54:39 -0700 Subject: [PATCH 1/3] Add MoE + expert-parallel Python API surface, PyTorch reference, and tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MoeEp exposes the constructor/forward contract for the fused SwiGLU MoE with EP dispatch (bf16, mxfp8, and nvfp4 public outputs); forward allocates the output representation until the device kernel lands. The pure-PyTorch MoeEpReference implements the full semantics — routing, variable-size all-to-all dispatch, local experts, Form-A top-k combine, and block-scaled quantization — and is validated by test/python/fe_api/moe_ep against naive per-token references, including 2-rank (gloo) and 4-GPU (NCCL) expert-parallel runs. Co-Authored-By: Claude Fable 5 --- docs/fe-oss-apis/moe_ep.md | 448 ++++++++++ docs/fe-oss-apis/overview.md | 1 + python/cudnn/__init__.py | 3 + python/cudnn/moe_ep/__init__.py | 6 + python/cudnn/moe_ep/api.py | 339 ++++++++ test/python/fe_api/moe_ep/moe_ep_reference.py | 792 ++++++++++++++++++ test/python/fe_api/moe_ep/test_moe_ep.py | 771 +++++++++++++++++ 7 files changed, 2360 insertions(+) create mode 100644 docs/fe-oss-apis/moe_ep.md create mode 100644 python/cudnn/moe_ep/__init__.py create mode 100644 python/cudnn/moe_ep/api.py create mode 100644 test/python/fe_api/moe_ep/moe_ep_reference.py create mode 100644 test/python/fe_api/moe_ep/test_moe_ep.py diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md new file mode 100644 index 000000000..7d2b57d55 --- /dev/null +++ b/docs/fe-oss-apis/moe_ep.md @@ -0,0 +1,448 @@ +# MoE + Expert Parallel API proposal + +Status: public API stub, implementation proposal, and executable PyTorch +reference. The API currently allocates uninitialized output storage without +launching a device kernel. Its numerical comparison is therefore marked as a +strict expected failure under `test/python/fe_api/moe_ep`; remove that marker +when backend execution is connected. + +This design is based on `mxfp8_glu_mega_instro.html`. It preserves that +kernel's routing, contiguous expert placement, SwiGLU, optional early router +weight, internal per-top-k combine plane, and final top-k reduction. It removes +workspace pointers, peer pointer mappers, streams, and scheduler tuning from the +semantic user interface. + +## Decision summary + +- The constructor contains static model, EP, capacity, and numerical choices. +- `__call__` contains only runtime tensors. +- Each rank supplies local tokens and local expert weights. Expert IDs in the + routing table are global. +- Expert ownership is contiguous and uniform. EP rank `r` owns + `[r * E_local, (r + 1) * E_local)`. +- `-1` is the only dropped/unused route value. Other out-of-range IDs are + errors. +- The first half of FC1 is `gate`; the second half is `up`. SwiGLU is + `silu(gate) * up`. +- `output_format` means the public, post-top-k-reduction `(T, H)` result. This is + an extension of the HTML interface, whose public result is BF16. +- `combine_format` independently describes each per-route FC2 contribution on + the EP return path. This corresponds to the HTML's internal `combine_quant` + and `combine_sf` planes. +- BF16, MXFP8, and NVFP4 are supported for both choices. Quantized output is a + data-plus-scale object; it is never represented as a scale-free PyTorch + tensor. + +The distinction between public output and combine traffic is intentional. If +"MXFP8/NVFP4 output" is meant only as a transport optimization, set +`combine_format` to that format and keep `output_format="bf16"`. + +## Proposed API + +The production class should be named `MoeEp`. The checked-in semantic +implementation is named `MoeEpReference` so it cannot be confused with a fused +kernel. + +```python +class MoeEp: + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[torch.distributed.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Literal["bf16", "mxfp8", "nvfp4"] = "bf16", + combine_format: Literal["bf16", "mxfp8", "nvfp4"] = "bf16", + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + ) -> None: ... + + def __call__( + self, + activation: Tensor | BlockScaledTensor, + fc1_weight: Tensor | BlockScaledTensor, + fc2_weight: Tensor | BlockScaledTensor, + topk_idx: Tensor, + topk_weights: Tensor, + ) -> ( + Tensor + | BlockScaledTensor + | tuple[Tensor | BlockScaledTensor, Tensor, Tensor] + ): ... +``` + +An initialized `ep_group` enables EP. `None` deliberately means a one-rank +execution, even if a default distributed process group exists. This prevents an +operator from silently communicating on the wrong group. + +### Constructor contract + +| Argument | Meaning | +|---|---| +| `num_experts` | Global expert count `E`; must be divisible by EP size. | +| `hidden_size` | Model hidden dimension `H`. | +| `intermediate_size` | Post-SwiGLU dimension `I`; FC1 has `2 * I` columns. | +| `top_k` | Fixed routing width `K`, with `1 <= K <= E`. | +| `ep_group` | Process group whose group-relative rank determines expert ownership. | +| `max_tokens_per_rank` | Maximum local input tokens `T`; optional in the reference and required by a static-workspace implementation. | +| `output_format` | Encoding returned after top-k reduction. | +| `combine_format` | Encoding/rounding of each route contribution before top-k reduction. | +| `apply_topk_in_fc1` | Multiply the post-SwiGLU intermediate by the router weight before FC2; otherwise multiply FC2 output. | +| `gate_up_clamp` | If set, use `gate = min(gate, abs(limit))` and `up = clamp(up, -abs(limit), abs(limit))`. | +| `generate_c` | Training integration: additionally return `fc1_c` (raw pre-SwiGLU FC1 accumulator of every route this rank's experts processed) and its row-aligned `route_metadata`. | + +The constructor validates static format alignment. MXFP8 output requires +`H % 32 == 0`; NVFP4 output requires `H % 16 == 0`. A production implementation +may also validate SM architecture, NVSHMEM availability, symmetric heap size, +and CUDA graph constraints here. + +Scheduler settings such as `token_back_mode`, `group_hint`, `flag_batch`, +`epi_flag_batch`, and scheduler stages belong in an optional backend/tuning +object. They do not change the result and should not be positional API +arguments. + +`fc2_in_kernel_topk_reduce` is different: BF16 atomic reduction can change the +rounding order. The first contract therefore specifies Form A semantics +(round each route to `combine_format`, reduce in FP32, encode once). Form B may +be used behind the API only if it satisfies the agreed tolerance; otherwise it +must become an explicit numerical mode rather than an invisible tuning flag. + +### Forward tensor contract + +Let `T` be this rank's token count and `E_local = E / ep_size`. + +| Tensor | Logical shape | Required properties | +|---|---:|---| +| `activation` | `(T, H)` | BF16/FP tensor, or block-scaled along axis 1. | +| `fc1_weight` | `(E_local, H, 2I)` | Local experts only; block-scaled weights use axis 1. | +| `fc2_weight` | `(E_local, I, H)` | Local experts only; block-scaled weights use axis 1. | +| `topk_idx` | `(T, K)` | INT32 or INT64 global expert IDs; `-1` means unused. | +| `topk_weights` | `(T, K)` | Floating router/combine weights. They are not implicitly normalized. | + +All tensors must be on one device. Quantized data and its scale tensor must +also share a device. Biases, shared experts, nonuniform expert placement, and +capacity-based route dropping are outside this first API. + +### Return value + +- BF16: a `torch.bfloat16` tensor with logical shape `(T, H)`. +- MXFP8/NVFP4: a `BlockScaledTensor` containing `data`, `scale`, `format`, + `logical_shape`, and the scaled axis. `dequantize()` reconstructs a regular + tensor. + +The return type is fixed by the constructor, so an individual module instance +does not change its output structure across calls. + +### Example + +```python +moe = MoeEpReference( + num_experts=8, + hidden_size=4096, + intermediate_size=14336, + top_k=2, + ep_group=ep_group, + max_tokens_per_rank=2048, + combine_format="mxfp8", + output_format="nvfp4", +) + +# Each rank passes its own tokens and its contiguous E_local weight shard. +output = moe(activation, local_fc1, local_fc2, topk_idx, topk_weights) +assert output.logical_shape == (activation.shape[0], 4096) +output_bf16 = output.dequantize(torch.bfloat16) +``` + +## Mathematical semantics + +For valid route `(t, k)` with global expert `e` and router weight `p[t, k]`: + +```text +z[t,k] = fp32(x[t]) @ fp32(W1[e]) +gate, up = split(z[t,k], I) +hidden[t,k] = silu(gate) * up + +if apply_topk_in_fc1: + hidden[t,k] *= p[t,k] + +expert[t,k] = hidden[t,k] @ fp32(W2[e]) + +if not apply_topk_in_fc1: + expert[t,k] *= p[t,k] + +combine[t,k] = dequantize(quantize(expert[t,k], combine_format)) +result[t] = sum_k(combine[t,k]) +output = encode(result, output_format) +``` + +For BF16 combine, `quantize/dequantize` above means a BF16 round trip. Invalid +slots contribute exact zero. Accumulation across top-k slots is FP32 and the +public encoding is applied after reduction. + +Moving the router weight across FC2 is algebraically equivalent in exact +arithmetic, but not necessarily after low-precision rounding. The option is +therefore semantic and is fixed in the constructor, matching the HTML kernel. + +## Block-scaled representation + +The API uses logical, unswizzled scales. Backend-specific F8_128x4 swizzling is +an implementation detail performed while constructing tensor maps or staging +weights. + +| Format | Payload | Scale | Block | Quantized axis | +|---|---|---|---:|---| +| BF16 | BF16 | none | n/a | n/a | +| MXFP8 | FP8 E4M3 | FP8 E8M0 | 32 | reduction/output axis | +| NVFP4 | packed FP4 E2M1, low nibble first | FP8 E4M3 | 16 | reduction/output axis | + +MXFP8 scale calculation is: + +```text +raw_scale = amax(block) / 448 +scale = 2 ** ceil(log2(raw_scale)) # E8M0 round toward +infinity +data = e4m3(clamp(block / scale, -448, 448)) +``` + +NVFP4 scale calculation is: + +```text +raw_scale = amax(block) / 6 +scale = e4m3(raw_scale) # round to nearest, saturate finite +data = e2m1(clamp(block / scale, -6, 6)) # two values per byte +``` + +E2M1 conversion in the reference uses round-to-nearest, ties-to-even. Logical +shapes may be padded to a complete block internally, but padding is not visible +through `logical_shape`. + +Examples of scale shapes are: + +| Logical tensor | MXFP8 scale | NVFP4 scale | +|---|---:|---:| +| activation `(T, H)` | `(T, ceil(H/32))` | `(T, ceil(H/16))` | +| FC1 `(E_local, H, 2I)` | `(E_local, ceil(H/32), 2I)` | `(E_local, ceil(H/16), 2I)` | +| FC2 `(E_local, I, H)` | `(E_local, ceil(I/32), H)` | `(E_local, ceil(I/16), H)` | +| output `(T, H)` | `(T, ceil(H/32))` | `(T, ceil(H/16))` | + +NVFP4 payload shape replaces the quantized axis by `ceil(axis_extent / 2)`. + +## Expert-parallel execution + +The semantic data flow is: + +```text +local x, topk_idx, topk_weights + | + v +flatten valid routes -> stable sort by destination EP rank + | + v +variable all-to-all dispatch (token, local expert id, route weight) + | + v +group by local expert -> FC1 -> SwiGLU -> FC2 -> combine-format round trip + | + v +reverse variable all-to-all in the exact dispatch order + | + v +scatter to local [token, top-k, hidden] plane -> FP32 top-k sum + | + v +encode public output +``` + +The reference uses two variable-split `all_to_all_single` phases. The target +MegaMoE kernel can use NVSHMEM pull for dispatch and direct remote stores or +token-back warps for return. Those are different transport mechanisms with the +same observable mapping. + +Stable ordering is required only so the reverse exchange can return results +without sending source token metadata to the expert rank. The source rank keeps +its local `(token, top-k slot)` permutation and scatters returned rows back into +the combine plane. + +Ranks may have different `T`. Zero-token ranks and zero-count peer splits must +participate in all collectives. A production workspace must reserve enough +inbound assignments for its documented capacity policy. The conservative bound +is `ep_size * max_tokens_per_rank * top_k`; a smaller bound requires an explicit +router capacity/drop contract. + +## Mapping to the MegaMoE HTML interface + +| HTML concept | Proposed API | +|---|---| +| `static_expert_shape=(E, 2I, H)` | `num_experts`, `intermediate_size`, `hidden_size` | +| `world_size` | inferred from `ep_group` | +| `num_topk` | `top_k` | +| `max_tokens_per_rank` | same name | +| `activation` + `activation_sf` | one `BlockScaledTensor` | +| `fc1_weight` + `fc1_weight_sf` | one `BlockScaledTensor` | +| `fc2_weight` + `fc2_weight_sf` | one `BlockScaledTensor` | +| `topk_idx`, `topk_weights` | same runtime tensors | +| internal `combine_quant`, `combine_sf` | selected by `combine_format`, not passed by the caller | +| BF16 `output_activation` | BF16 case of the returned value | +| new quantized public output | `BlockScaledTensor` selected by `output_format` | +| `local_workspace`, `shared_workspace` | owned/cached by the implementation | +| `peer_rank_ptr_mapper_host` | derived from the EP communication backend | +| `max_active_clusters`, `stream` | backend launch state; current PyTorch stream is used | +| `token_comm_args` | private lowering/kernel argument bundle | +| `generate_c`, `fc1_c` | same names; `fc1_c` is the second returned value | +| `src_token_topk_idx`, `token_src_metadata` | `route_metadata`, the third returned value | + +### `fc1_c` and `route_metadata` (training integration) + +With `generate_c=True`, `__call__` returns `(output, fc1_c, route_metadata)`. +`fc1_c` is BF16 with shape `(local_routes, 2 * intermediate_size)`, where +`local_routes` is the data-dependent number of valid routes assigned to this +rank's experts. It stays **expert-rank-local** — matching the kernel, which +writes `fc1_c` where FC1 ran and never ships it back — because the backward +pass re-dispatches gradients to the expert ranks, which is where the stashed +activations are consumed. + +Row semantics: grouped by local expert (ascending); within an expert, ordered +by source rank, then the source rank's token-major route order. Values are the +raw FC1 accumulator captured **before** SwiGLU, before the gate/up clamp, and +without the router weight (which applies after SwiGLU when +`apply_topk_in_fc1=True`). The kernel's 128-row per-expert padding is a layout +detail and is absent from the logical contract. + +`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`. This is the information the backward pass needs to +re-dispatch output gradients to the right expert-rank rows and to scatter +input gradients back to `(src_token, src_slot)` on the source ranks. It is +the public form of the kernel's `src_token_topk_idx`/`token_src_metadata` +routing words, which the dispatch phase already materializes on the expert +rank. + +### Saved tensors for backward + +`MoeEpReference.backward(grad_output, activation, fc1_weight, fc2_weight, +topk_idx, topk_weights, fc1_c, route_metadata)` returns +`(grad_activation, grad_fc1_weight, grad_fc2_weight, grad_topk_weights)` and +is the executable statement of the save-set. What must survive from forward +to backward: + +| Tensor | Where it lives | Why backward needs it | +|---|---|---| +| `fc1_c` | expert rank (stash) | Sole recompute source: gate/up split, clamp masks, SwiGLU, and the FC2 input `h` are rebuilt from it. Also yields `dW1 = xᵀ · d_c` and `d_x = d_c · W1ᵀ`. | +| `route_metadata` | expert rank (stash) | Reconstructs the receive-order ↔ `fc1_c`-row permutation (sort by `(src_rank, src_token, src_slot)`), groups rows by local expert, and drives the gradient return scatter. | +| `fc1_weight`, `fc2_weight` | expert rank (resident params) | `d_x = d_c · W1ᵀ`, `d_h = d_y · W2ᵀ`. | +| `activation` | source rank (framework input) | FC1 input `x` for `dW1`. Re-dispatched in backward along the identical routes (`topk_idx` is deterministic), or alternatively the forward's dispatched copy `(pool, H)` is stashed on the expert rank to trade memory for the second dispatch. | +| `topk_idx`, `topk_weights` | source rank (framework inputs) | `topk_idx` regenerates the exact dispatch plan; `topk_weights` reconstructs `h' = w·h` for `dW2` and produces `d_w` (returned per `(src_token, src_slot)`). | +| `grad_output` | source rank (incoming) | One row per route is dispatched to the expert rank: `d_y[route] = grad_output[src_token]` (top-k reduce is a sum). | + +Deliberately **not** saved: the post-SwiGLU `fc1_output`/`fc1_output_sf` +(recomputed from `fc1_c`), the `combine_quant`/`combine_sf` planes, the +public output, and all counters/flags. + +Numerical conventions of the reference backward, to be matched by a fused +kernel: all quantization round-trips (input decode, `combine_format`, +`output_format`) are straight-through — `grad_output` is the gradient of the +dequantized `(T, H)` output; the bf16 `fc1_c` stash is the recompute source, +so backward SwiGLU math runs on bf16-rounded accumulator values (standard +stash precision); clamp gradients are inclusive at the bounds, matching +`torch.clamp`; when `apply_topk_in_fc1=True` the router-weight gradient is +`d_w = ⟨d_h', h⟩` with `h` pre-weight, otherwise `d_w = ⟨d_y, y_pre⟩`. + +## Production implementation plan + +1. **Constructor and plan cache** + - Resolve group-relative EP rank/size and local expert range. + - Validate static dimensions, format combinations, architecture, and capacity. + - Build a cache key from static configuration plus runtime data/scale dtypes + and strides. + - Size local and symmetric workspaces. Allocate through a backend context, + not on every call. + +2. **Forward validation and views** + - Flatten `BlockScaledTensor` objects into payload/scale kernel arguments. + - Validate `(T, K)` routing and local weight descriptors. + - Slice internal dispatch, counter, FC1, combine, and scale planes from owned + workspaces. The caller never passes these pointers. + +3. **Dispatch** + - Count valid routes per destination and expert. + - Prefix-sum counts, place routing words, and transfer activation payload, + activation scale, and router weight. + - Preserve `(source rank, source token, source slot)` in compact metadata or + in a reversible placement order. + +4. **Local expert kernel** + - Use one specialization per input/weight family and combine format. + - Accumulate GEMMs in FP32. + - Apply the documented gate/up clamp and router-weight location. + - Quantize route contributions with block boundaries aligned to `H`. + +5. **Return and reduction** + - Direct epilogue remote stores are the default fast path. + - Standalone/reused token-back warps remain tuning choices. + - Reduce the internal top-k plane in FP32. + - BF16-cast or block-quantize the reduced result according to + `output_format`. + +6. **Runtime behavior** + - Use the current PyTorch CUDA stream. + - Avoid host synchronization after plan creation. + - Reset counters in-kernel so repeated calls and CUDA graph replay are safe. + - Keep peer-visible allocation addresses stable for the lifetime of the + plan/workspace object. + +For quantized public output, the fused final top-k reducer should emit payload +and logical scales together. It should not first materialize a BF16 `(T, H)` +buffer unless that fallback is selected. + +## Reference implementation + +The executable reference is +`test/python/fe_api/moe_ep/moe_ep_reference.py`. It accepts ordinary floating +tensors or `BlockScaledTensor` inputs and weights. With an explicit process +group it executes actual variable-size PyTorch collectives, so it checks both +MoE math and EP ordering. + +The reference is a correctness oracle, not a performance model: + +- GEMMs and top-k accumulation use FP32. +- It models combine and public-output rounding, but not CTA tile-dependent + accumulation order. +- Scale tensors are logical, not atom-swizzled. +- It uses collective push communication rather than NVSHMEM pull/remote store. +- Distributed PyTorch collectives in this reference are not an autograd + implementation. The initial scope is inference forward. +- It has no expert-capacity drop policy beyond explicit `-1` routes. + +## Validation and test matrix + +The accompanying tests cover: + +- MXFP8 and NVFP4 quantize/dequantize shape and dtype contracts; +- quantization along the weight reduction axis; +- BF16 output for both router-weight locations and gate/up clamp; +- MXFP8 and NVFP4 public outputs; +- mixed block-scaled activation, FC1, and FC2 inputs; +- invalid expert IDs; +- two EP ranks with unequal local token counts and cross-rank routes. + +Production validation should add CUDA/NVSHMEM runs for every pair of +`combine_format` and `output_format`, skewed/all-to-one routing, empty experts, +zero-token ranks, duplicate expert selections, maximum capacity, CUDA graph +replay, and comparisons against this reference after dequantization. + +## Proposed first-version boundaries + +These choices should remain explicit until there is a concrete model requiring +more surface area: + +- contiguous, equal expert partition only; +- no expert bias; +- no shared/dense expert inside this operator; +- no implicit top-k normalization; +- no capacity factor or implicit route drop; +- no backward API; +- no auxiliary FC1 capture; +- logical scales at the Python boundary, backend swizzle internally. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 2bc3ccfb2..5293df237 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -26,6 +26,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [SDPA Forward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-forward-fe-oss-sm100-d256) - [SDPA Backward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-backward-fe-oss-sm100-d256) - [RMSNorm + SiLU](rmsnorm_silu.md) +- [MoE + Expert Parallel API proposal](moe_ep.md) ## Installation and setup diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 6e8b99829..c06698b95 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -278,6 +278,9 @@ def _dlopen_cudnn(): _OPTIONAL_DEPENDENCY_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" _LAZY_OPTIONAL_IMPORTS = { + "moe_ep": (".moe_ep", None), + "MoeEp": (".moe_ep", "MoeEp"), + "MoeFormat": (".moe_ep", "MoeFormat"), "BSA": (".block_sparse_attention", "BSA"), "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), "block_sparse_attention_backward": (".block_sparse_attention", "block_sparse_attention_backward"), diff --git a/python/cudnn/moe_ep/__init__.py b/python/cudnn/moe_ep/__init__.py new file mode 100644 index 000000000..ad1f24db5 --- /dev/null +++ b/python/cudnn/moe_ep/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from .api import BlockScaledTensor, MoeEp, MoeFormat, MoeTensor + +__all__ = ["BlockScaledTensor", "MoeEp", "MoeFormat", "MoeTensor"] diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py new file mode 100644 index 000000000..45d6b9cd1 --- /dev/null +++ b/python/cudnn/moe_ep/api.py @@ -0,0 +1,339 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Python API surface for fused SwiGLU MoE with expert parallelism. + +The public contract is present before the device implementation so callers and +the PyTorch reference can be wired into tests. ``MoeEp.__call__`` currently +allocates output storage but does not launch a backend kernel. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Tuple, Union + +import torch +import torch.distributed as dist + + +class MoeFormat(str, Enum): + """Data formats supported by the MoE+EP interface.""" + + 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 + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Data-plus-scale result returned for MXFP8 and NVFP4 outputs.""" + + 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") + logical_shape = tuple(int(dim) for dim in self.logical_shape) + axis = _normalize_axis(self.axis, len(logical_shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", logical_shape) + object.__setattr__(self, "axis", axis) + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Decode the logical, unswizzled block-scaled representation.""" + + 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.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) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +def _logical_shape(tensor: MoeTensor) -> Tuple[int, ...]: + if isinstance(tensor, BlockScaledTensor): + return tensor.logical_shape + return tuple(tensor.shape) + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +class MoeEp: + """Fused SwiGLU MoE operator with contiguous expert parallel sharding. + + Global expert ``e`` belongs to group-relative EP rank + ``e // experts_per_rank``. The constructor captures static configuration; + calling the instance accepts runtime tensors for this rank. + + With ``generate_c=True`` (training integration), ``__call__`` additionally + returns ``fc1_c`` and ``route_metadata``. ``fc1_c`` is the raw pre-SwiGLU + FC1 accumulator for every route this rank's experts processed, BF16, shape + ``(local_routes, 2 * intermediate)``. Rows are grouped by local expert + (ascending) and ordered within each expert by source rank, then the source + rank's token-major route order. The rows are captured before the gate/up + clamp and carry no router weight. ``route_metadata`` is Int32 + ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``, row-aligned with + ``fc1_c``, identifying each route for the backward gradient re-dispatch. + + The backend launch is intentionally a TODO. Until it is implemented, + ``__call__`` returns newly allocated, uninitialized output storage with the + correct public representation. + """ + + 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, + ) -> 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 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) + + 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 " f"{required_multiple} for {name}={fmt.value}") + + def _allocate_output(self, token_count: int, device: torch.device) -> MoeTensor: + logical_shape = (token_count, self.hidden_size) + if self.output_format is MoeFormat.BF16: + return torch.empty(logical_shape, dtype=torch.bfloat16, device=device) + if self.output_format is MoeFormat.MXFP8: + data = torch.empty( + logical_shape, + dtype=_require_torch_dtype("float8_e4m3fn"), + device=device, + ) + scale = torch.empty( + (token_count, self.hidden_size // 32), + dtype=_require_torch_dtype("float8_e8m0fnu"), + device=device, + ) + else: + data = torch.empty( + (token_count, self.hidden_size // 2), + dtype=torch.uint8, + device=device, + ) + scale = torch.empty( + (token_count, self.hidden_size // 16), + dtype=_require_torch_dtype("float8_e4m3fn"), + device=device, + ) + return BlockScaledTensor( + data=data, + scale=scale, + format=self.output_format, + logical_shape=logical_shape, + axis=-1, + ) + + def _count_local_routes(self, topk_idx: torch.Tensor) -> int: + """Number of valid routes this rank's experts receive. + + Data-dependent: single-rank counts locally, EP exchanges per-rank + route counts (the same exchange the device dispatch performs). + """ + + flat = topk_idx.reshape(-1).to(torch.int64) + expert = flat[flat != -1] + if expert.numel() > 0 and bool(((expert < 0) | (expert >= self.num_experts)).any().item()): + raise ValueError("topk_idx contains out-of-range expert ids") + if self.ep_size == 1: + return int(expert.numel()) + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + send_counts = torch.bincount(destination, minlength=self.ep_size) + if send_counts.device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + send_counts = send_counts.cpu() + recv_counts = torch.empty_like(send_counts) + dist.all_to_all_single(recv_counts, send_counts, group=self.ep_group) + return int(recv_counts.sum().item()) + + 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]]: + """Allocate the result for a future backend MoE+EP launch. + + Expected logical shapes are ``activation=(T,H)``, + ``fc1_weight=(E_local,H,2I)``, ``fc2_weight=(E_local,I,H)``, and + ``topk_idx=topk_weights=(T,K)``. + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when the operator was constructed with ``generate_c=True``. + """ + + activation_shape = _logical_shape(activation) + if len(activation_shape) != 2 or activation_shape[1] != self.hidden_size: + raise ValueError(f"activation logical shape must be (T, {self.hidden_size}), got {activation_shape}") + token_count = activation_shape[0] + expected_fc1 = (self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size) + expected_fc2 = (self.experts_per_rank, self.intermediate_size, self.hidden_size) + if _logical_shape(fc1_weight) != expected_fc1: + raise ValueError(f"fc1_weight logical shape must be {expected_fc1}") + if _logical_shape(fc2_weight) != expected_fc2: + raise ValueError(f"fc2_weight logical shape must be {expected_fc2}") + 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 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) + for name, tensor in ( + ("fc1_weight", fc1_weight), + ("fc2_weight", fc2_weight), + ("topk_idx", topk_idx), + ("topk_weights", topk_weights), + ): + if _tensor_device(tensor) != device: + raise ValueError(f"{name} must be on {device}, got {_tensor_device(tensor)}") + + # TODO: dispatch routes, launch local expert FC1/SwiGLU/FC2, return + # contributions, reduce top-k, and write this allocation. + output = self._allocate_output(token_count, device) + if not self.generate_c: + return output + local_routes = self._count_local_routes(topk_idx) + fc1_c = torch.empty( + (local_routes, 2 * self.intermediate_size), + dtype=torch.bfloat16, + device=device, + ) + route_metadata = torch.empty((local_routes, 4), dtype=torch.int32, device=device) + return output, fc1_c, route_metadata + + +__all__ = ["BlockScaledTensor", "MoeEp", "MoeFormat", "MoeTensor"] diff --git a/test/python/fe_api/moe_ep/moe_ep_reference.py b/test/python/fe_api/moe_ep/moe_ep_reference.py new file mode 100644 index 000000000..3fab88a50 --- /dev/null +++ b/test/python/fe_api/moe_ep/moe_ep_reference.py @@ -0,0 +1,792 @@ +# Copyright (c) 2026 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, +) -> 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() + + 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.float() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=-1).dequantize() + + +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, + ) -> 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 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) + + 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 {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})" + ) + + 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=torch.float32, + 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.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up + weights = route_weight.index_select(0, positions).unsqueeze(-1) + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + expert_output = intermediate @ fc2_weight[expert] + if not self.apply_topk_in_fc1: + expert_output = expert_output * weights + expert_output = _format_round_trip(expert_output, self.combine_format) + 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 = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_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) + + 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_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) + 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)`` in 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 {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 = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_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 {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.float() + recv_tokens = self._all_to_all(activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + 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.float() + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, 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) + 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.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate, up + sig = torch.sigmoid(g) + s = g * sig + h = s * u + + if self.apply_topk_in_fc1: + h_fc2 = h * w + d_y_pre = d_y + else: + h_fc2 = h + d_y_pre = d_y * w + grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + 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 * (h @ fc2_float[expert])).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 <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1) + 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) + 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/test/python/fe_api/moe_ep/test_moe_ep.py b/test/python/fe_api/moe_ep/test_moe_ep.py new file mode 100644 index 000000000..083b46f5a --- /dev/null +++ b/test/python/fe_api/moe_ep/test_moe_ep.py @@ -0,0 +1,771 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from fe_api.moe_ep.moe_ep_reference import ( + BlockScaledTensor, + MoeEpReference, + MoeFormat, + quantize_blockwise, +) + +pytestmark = pytest.mark.L0 + + +def _output_as_float(output): + if isinstance(output, torch.Tensor): + return output.float() + return output.dequantize() + + +def _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + *, + apply_topk_in_fc1, + clamp=None, + combine_format=MoeFormat.BF16, +): + token_count, top_k = topk_idx.shape + hidden_size = activation.shape[1] + intermediate_size = fc2_weight.shape[1] + combine = torch.zeros(token_count, top_k, hidden_size, dtype=torch.float32) + for token in range(token_count): + for slot in range(top_k): + expert = int(topk_idx[token, slot]) + if expert == -1: + continue + gate_up = activation[token].float() @ fc1_weight[expert].float() + gate, up = gate_up.split(intermediate_size) + if clamp is not None: + gate = gate.clamp(max=clamp) + up = up.clamp(-clamp, clamp) + intermediate = F.silu(gate) * up + route_weight = topk_weights[token, slot].float() + if apply_topk_in_fc1: + intermediate = intermediate * route_weight + result = intermediate @ fc2_weight[expert].float() + if not apply_topk_in_fc1: + result = result * route_weight + if combine_format is MoeFormat.BF16: + result = result.to(torch.bfloat16).float() + else: + result = quantize_blockwise(result, combine_format).dequantize() + combine[token, slot] = result + return combine.sum(dim=1).to(torch.bfloat16) + + +@pytest.mark.parametrize( + "format,expected_data_shape,expected_scale_shape,scale_dtype", + [ + (MoeFormat.MXFP8, (3, 64), (3, 2), torch.float8_e8m0fnu), + (MoeFormat.NVFP4, (3, 32), (3, 4), torch.float8_e4m3fn), + ], +) +def test_block_scaled_round_trip(format, expected_data_shape, expected_scale_shape, scale_dtype): + values = torch.linspace(-4.0, 4.0, 3 * 64).reshape(3, 64) + quantized = quantize_blockwise(values, format) + + assert isinstance(quantized, BlockScaledTensor) + assert quantized.format is format + assert quantized.logical_shape == (3, 64) + assert tuple(quantized.data.shape) == expected_data_shape + assert tuple(quantized.scale.shape) == expected_scale_shape + assert quantized.scale.dtype == scale_dtype + assert quantized.dequantize().shape == values.shape + assert torch.isfinite(quantized.dequantize()).all() + + +def test_quantization_along_weight_reduction_axis(): + values = torch.randn(2, 64, 48) + quantized = quantize_blockwise(values, MoeFormat.NVFP4, axis=1) + + assert quantized.logical_shape == (2, 64, 48) + assert quantized.data.shape == (2, 32, 48) + assert quantized.scale.shape == (2, 4, 48) + assert quantized.dequantize().shape == values.shape + + +@pytest.mark.parametrize("apply_topk_in_fc1", [False, True]) +def test_single_rank_bf16_matches_naive(apply_topk_in_fc1): + torch.manual_seed(7) + experts, tokens, hidden, intermediate, top_k = 4, 5, 32, 24, 2 + activation = torch.randn(tokens, hidden, dtype=torch.bfloat16) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate, dtype=torch.bfloat16) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden, dtype=torch.bfloat16) / 8 + topk_idx = torch.tensor([[0, 3], [2, 1], [1, -1], [3, 0], [2, 0]], dtype=torch.int64) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.2, 0.8], [1.0, 0.0], [0.55, 0.45], [0.6, 0.4]], + dtype=torch.float32, + ) + + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + combine_format="bf16", + output_format="bf16", + apply_topk_in_fc1=apply_topk_in_fc1, + gate_up_clamp=2.5, + ) + actual = op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + expected = _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + apply_topk_in_fc1=apply_topk_in_fc1, + clamp=2.5, + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_generate_c_single_rank_matches_naive(): + """fc1_c is the pre-clamp, unweighted gate+up accumulator, grouped by expert.""" + + torch.manual_seed(37) + experts, tokens, hidden, intermediate, top_k = 4, 5, 32, 24, 2 + activation = torch.randn(tokens, hidden) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden) / 8 + topk_idx = torch.tensor([[0, 3], [2, 1], [1, -1], [3, 0], [2, 0]], dtype=torch.int64) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.2, 0.8], [1.0, 0.0], [0.55, 0.45], [0.6, 0.4]], + dtype=torch.float32, + ) + + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + gate_up_clamp=2.5, + generate_c=True, + ) + output, fc1_c, route_metadata = op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + + expected_rows = [] + expected_metadata = [] + for expert in range(experts): + for token in range(tokens): + for slot in range(top_k): + if int(topk_idx[token, slot]) == expert: + expected_rows.append(activation[token] @ fc1_weight[expert]) + expected_metadata.append([expert, 0, token, slot]) + expected_fc1_c = torch.stack(expected_rows).to(torch.bfloat16) + + assert fc1_c.dtype == torch.bfloat16 + assert fc1_c.shape == (int((topk_idx != -1).sum()), 2 * intermediate) + torch.testing.assert_close(fc1_c, expected_fc1_c, atol=0, rtol=0) + + assert route_metadata.dtype == torch.int32 + assert route_metadata.tolist() == expected_metadata + + expected_output = _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + apply_topk_in_fc1=True, + clamp=2.5, + ) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + + +def _autograd_replica_grads( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + grad_output, + *, + apply_topk_in_fc1=True, + clamp=None, +): + """Autograd gradients of a per-route replica of the reference forward. + + The FC1 accumulator is rounded to bf16 with a straight-through gradient, + matching a backward pass that recomputes SwiGLU from the bf16 fc1_c stash. + Combine/output round-trips are omitted (straight-through), as in + ``MoeEpReference.backward``. + """ + + a = activation.detach().clone().float().requires_grad_() + w1 = fc1_weight.detach().clone().float().requires_grad_() + w2 = fc2_weight.detach().clone().float().requires_grad_() + tw = topk_weights.detach().clone().float().requires_grad_() + token_count, top_k = topk_idx.shape + intermediate = w2.shape[1] + rows = [] + for token in range(token_count): + acc = torch.zeros(a.shape[1], dtype=torch.float32) + for slot in range(top_k): + expert = int(topk_idx[token, slot]) + if expert == -1: + continue + c = a[token] @ w1[expert] + c = c + (c.to(torch.bfloat16).float() - c).detach() + gate, up = c.split(intermediate) + if clamp is not None: + gate = gate.clamp(max=clamp) + up = up.clamp(-clamp, clamp) + h = F.silu(gate) * up + weight = tw[token, slot] + if apply_topk_in_fc1: + h = h * weight + y = h @ w2[expert] + if not apply_topk_in_fc1: + y = y * weight + acc = acc + y + rows.append(acc) + torch.stack(rows).backward(grad_output.float()) + return a.grad, w1.grad, w2.grad, tw.grad + + +@pytest.mark.parametrize("apply_topk_in_fc1", [False, True]) +def test_backward_single_rank_matches_autograd(apply_topk_in_fc1): + torch.manual_seed(43) + experts, tokens, hidden, intermediate, top_k = 4, 5, 32, 24, 2 + activation = torch.randn(tokens, hidden) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden) / 8 + topk_idx = torch.tensor([[0, 3], [2, 1], [1, -1], [3, 0], [2, 0]], dtype=torch.int64) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.2, 0.8], [1.0, 0.0], [0.55, 0.45], [0.6, 0.4]], + dtype=torch.float32, + ) + grad_output = torch.randn(tokens, hidden) + + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + apply_topk_in_fc1=apply_topk_in_fc1, + gate_up_clamp=2.5, + generate_c=True, + ) + _, fc1_c, route_metadata = op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + grad_activation, grad_fc1, grad_fc2, grad_topk = op.backward( + grad_output, + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + expected_ga, expected_g1, expected_g2, expected_gw = _autograd_replica_grads( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + grad_output, + apply_topk_in_fc1=apply_topk_in_fc1, + clamp=2.5, + ) + + torch.testing.assert_close(grad_activation, expected_ga, rtol=1e-4, atol=1e-6) + torch.testing.assert_close(grad_fc1, expected_g1, rtol=1e-4, atol=1e-6) + torch.testing.assert_close(grad_fc2, expected_g2, rtol=1e-4, atol=1e-6) + torch.testing.assert_close(grad_topk, expected_gw, rtol=1e-4, atol=1e-6) + assert grad_topk[2, 1].item() == 0.0 # -1 slot receives no gradient + + +@pytest.mark.parametrize("input_format", [MoeFormat.MXFP8, MoeFormat.NVFP4]) +def test_backward_quantized_inputs_matches_autograd(input_format): + """Backward with block-scaled inputs and quantized combine/output formats. + + Gradients are straight-through with respect to the dequantized values, so + the expected gradients come from the autograd replica run on the decoded + tensors; the combine/output encodes must not perturb backward. + """ + + torch.manual_seed(47) + experts, tokens, hidden, intermediate, top_k = 2, 4, 32, 16, 2 + device = torch.device("cpu") + activation, fc1_weight, fc2_weight = _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device) + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64) + topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [1.0, 0.0], [0.55, 0.45]]) + grad_output = torch.randn(tokens, hidden) + + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + output_format=input_format, + combine_format=input_format, + generate_c=True, + ) + _, fc1_c, route_metadata = op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + grad_activation, grad_fc1, grad_fc2, grad_topk = op.backward( + grad_output, + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + expected_ga, expected_g1, expected_g2, expected_gw = _autograd_replica_grads( + activation.dequantize(), + fc1_weight.dequantize(), + fc2_weight.dequantize(), + topk_idx, + topk_weights, + grad_output, + ) + + torch.testing.assert_close(grad_activation, expected_ga, rtol=1e-4, atol=1e-6) + torch.testing.assert_close(grad_fc1, expected_g1, rtol=1e-4, atol=1e-6) + torch.testing.assert_close(grad_fc2, expected_g2, rtol=1e-4, atol=1e-6) + torch.testing.assert_close(grad_topk, expected_gw, rtol=1e-4, atol=1e-6) + + +@pytest.mark.parametrize("combine_format", [MoeFormat.MXFP8, MoeFormat.NVFP4]) +def test_quantized_combine_round_trip(combine_format): + torch.manual_seed(13) + experts, tokens, hidden, intermediate = 2, 4, 32, 16 + activation = torch.randn(tokens, hidden) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden) / 8 + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64) + topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [1.0, 0.0], [0.55, 0.45]]) + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=2, + combine_format=combine_format, + ) + + actual = op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + expected = _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + apply_topk_in_fc1=True, + combine_format=combine_format, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_quantized_inputs_and_weights_use_logical_shapes(): + torch.manual_seed(19) + experts, tokens, hidden, intermediate = 2, 3, 32, 32 + activation = torch.randn(tokens, hidden) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden) / 8 + q_activation = quantize_blockwise(activation, "mxfp8", axis=1) + q_fc1 = quantize_blockwise(fc1_weight, "mxfp8", axis=1) + q_fc2 = quantize_blockwise(fc2_weight, "nvfp4", axis=1) + topk_idx = torch.tensor([[0], [1], [0]], dtype=torch.int64) + topk_weights = torch.ones(tokens, 1) + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=1, + ) + + actual = op(q_activation, q_fc1, q_fc2, topk_idx, topk_weights) + expected = _naive_reference( + q_activation.dequantize(), + q_fc1.dequantize(), + q_fc2.dequantize(), + topk_idx, + topk_weights, + apply_topk_in_fc1=True, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("output_format", ["bf16", "mxfp8", "nvfp4"]) +@pytest.mark.xfail( + strict=True, + reason="MoeEp.__call__ only allocates output; the device implementation is not wired yet", +) +def test_moe_ep_api_matches_reference(output_format): + """Exercise the public API and reference through the same tensor contract.""" + + from cudnn import MoeEp + + torch.manual_seed(23) + device = torch.device("cuda") + experts, tokens, hidden, intermediate, top_k = 2, 4, 32, 16, 2 + activation = torch.randn(tokens, hidden, dtype=torch.bfloat16, device=device) + fc1_weight = ( + torch.randn( + experts, + hidden, + 2 * intermediate, + dtype=torch.bfloat16, + device=device, + ) + / 8 + ) + fc2_weight = ( + torch.randn( + experts, + intermediate, + hidden, + dtype=torch.bfloat16, + device=device, + ) + / 8 + ) + topk_idx = torch.tensor( + [[0, 1], [1, 0], [0, -1], [1, 0]], + dtype=torch.int64, + device=device, + ) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.6, 0.4], [1.0, 0.0], [0.55, 0.45]], + dtype=torch.float32, + device=device, + ) + + kwargs = dict( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + output_format=output_format, + combine_format="bf16", + ) + api = MoeEp(**kwargs) + reference = MoeEpReference(**kwargs) + + actual = api(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + expected = reference(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + + actual_shape = tuple(actual.shape) if isinstance(actual, torch.Tensor) else actual.logical_shape + assert actual_shape == (tokens, hidden) + torch.testing.assert_close( + _output_as_float(actual), + _output_as_float(expected), + atol=0, + rtol=0, + ) + + +def _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device): + """Quantized inputs whose scale factors ride inside each BlockScaledTensor. + + Everything is block-scaled along the GEMM reduction axis: activation along + hidden, fc1_weight along hidden, fc2_weight along intermediate. + """ + + activation = quantize_blockwise( + torch.randn(tokens, hidden, device=device), + input_format, + axis=1, + ) + fc1_weight = quantize_blockwise( + torch.randn(experts, hidden, 2 * intermediate, device=device) / 8, + input_format, + axis=1, + ) + fc2_weight = quantize_blockwise( + torch.randn(experts, intermediate, hidden, device=device) / 8, + input_format, + axis=1, + ) + return activation, fc1_weight, fc2_weight + + +@pytest.mark.parametrize("input_format", [MoeFormat.MXFP8, MoeFormat.NVFP4]) +def test_moe_ep_api_accepts_block_scaled_inputs(input_format): + """The API takes data+scale bundles where the kernel takes separate sf args.""" + + from cudnn import MoeEp + + torch.manual_seed(29) + device = torch.device("cuda") + experts, tokens, hidden, intermediate = 2, 4, 32, 16 + activation, fc1_weight, fc2_weight = _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device) + + api = MoeEp(num_experts=experts, hidden_size=hidden, intermediate_size=intermediate, top_k=1) + output = api( + activation, + fc1_weight, + fc2_weight, + torch.tensor([[0], [1], [0], [1]], dtype=torch.int64, device=device), + torch.ones(tokens, 1, device=device), + ) + + assert isinstance(output, torch.Tensor) + assert output.shape == (tokens, hidden) + assert output.dtype == torch.bfloat16 + + +def test_moe_ep_api_generate_c_allocates_fc1_c(): + """generate_c=True returns (output, fc1_c) sized by this rank's valid routes.""" + + from cudnn import MoeEp + + device = torch.device("cuda") + experts, tokens, hidden, intermediate = 2, 4, 32, 16 + api = MoeEp( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=2, + generate_c=True, + ) + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64, device=device) + output, fc1_c, route_metadata = api( + torch.empty(tokens, hidden, dtype=torch.bfloat16, device=device), + torch.empty(experts, hidden, 2 * intermediate, dtype=torch.bfloat16, device=device), + torch.empty(experts, intermediate, hidden, dtype=torch.bfloat16, device=device), + topk_idx, + torch.ones(tokens, 2, dtype=torch.float32, device=device), + ) + + valid_routes = int((topk_idx != -1).sum()) + assert output.shape == (tokens, hidden) + assert fc1_c.dtype == torch.bfloat16 + assert fc1_c.shape == (valid_routes, 2 * intermediate) + assert fc1_c.device.type == "cuda" + assert route_metadata.dtype == torch.int32 + assert route_metadata.shape == (valid_routes, 4) + assert route_metadata.device.type == "cuda" + + +@pytest.mark.xfail( + strict=True, + reason="MoeEp.__call__ only allocates output; the device implementation is not wired yet", +) +def test_moe_ep_api_generate_c_matches_reference(): + from cudnn import MoeEp + + torch.manual_seed(41) + device = torch.device("cuda") + experts, tokens, hidden, intermediate, top_k = 2, 4, 32, 16, 2 + activation = torch.randn(tokens, hidden, dtype=torch.bfloat16, device=device) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate, dtype=torch.bfloat16, device=device) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden, dtype=torch.bfloat16, device=device) / 8 + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64, device=device) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.6, 0.4], [1.0, 0.0], [0.55, 0.45]], + dtype=torch.float32, + device=device, + ) + + kwargs = dict( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + generate_c=True, + ) + actual, actual_fc1_c, actual_metadata = MoeEp(**kwargs)(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + expected, expected_fc1_c, expected_metadata = MoeEpReference(**kwargs)(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + + torch.testing.assert_close(_output_as_float(actual), _output_as_float(expected), atol=0, rtol=0) + torch.testing.assert_close(actual_fc1_c, expected_fc1_c, atol=0, rtol=0) + torch.testing.assert_close(actual_metadata, expected_metadata, atol=0, rtol=0) + + +@pytest.mark.parametrize("input_format", ["mxfp8", "nvfp4"]) +@pytest.mark.xfail( + strict=True, + reason="MoeEp.__call__ only allocates output; the device implementation is not wired yet", +) +def test_moe_ep_api_quantized_inputs_match_reference(input_format): + """Quantized activations/weights carry their scale factors through both paths.""" + + from cudnn import MoeEp + + torch.manual_seed(31) + device = torch.device("cuda") + experts, tokens, hidden, intermediate, top_k = 2, 4, 32, 16, 2 + activation, fc1_weight, fc2_weight = _block_scaled_inputs(input_format, experts, tokens, hidden, intermediate, device) + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64, device=device) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.6, 0.4], [1.0, 0.0], [0.55, 0.45]], + dtype=torch.float32, + device=device, + ) + + kwargs = dict(num_experts=experts, hidden_size=hidden, intermediate_size=intermediate, top_k=top_k) + actual = MoeEp(**kwargs)(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + expected = MoeEpReference(**kwargs)(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + + torch.testing.assert_close( + _output_as_float(actual), + _output_as_float(expected), + atol=0, + rtol=0, + ) + + +def _distributed_worker(rank, world_size, init_file): + device = torch.device("cuda", rank % torch.cuda.device_count()) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + experts, hidden, intermediate, top_k = 2 * world_size, 32, 16, 2 + torch.manual_seed(1234) + global_fc1 = torch.randn(experts, hidden, 2 * intermediate) / 8 + global_fc2 = torch.randn(experts, intermediate, hidden) / 8 + experts_per_rank = experts // world_size + begin = rank * experts_per_rank + end = begin + experts_per_rank + + generator = torch.Generator().manual_seed(2000 + rank) + token_count = 3 + rank + activation = torch.randn(token_count, hidden, generator=generator) + topk_idx = torch.tensor( + [[(rank + token) % experts, (rank + token + 2) % experts] for token in range(token_count)], + dtype=torch.int64, + ) + topk_weights = torch.tensor([[0.625, 0.375]]).expand(token_count, -1).contiguous() + + op = MoeEpReference( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + ep_group=dist.group.WORLD, + combine_format="bf16", + output_format="bf16", + generate_c=True, + ) + actual, fc1_c, route_metadata = op( + activation.to(device), + global_fc1[begin:end].contiguous().to(device), + global_fc2[begin:end].contiguous().to(device), + topk_idx.to(device), + topk_weights.to(device), + ) + assert actual.device.type == "cuda" + expected = _naive_reference( + activation, + global_fc1, + global_fc2, + topk_idx, + topk_weights, + apply_topk_in_fc1=True, + ) + # GPU and CPU float32 GEMMs accumulate in different orders, so the + # cross-device comparison uses bf16-level tolerances instead of exact. + torch.testing.assert_close(actual.cpu(), expected) + + # fc1_c contract: rows grouped by local expert, ordered by source rank, + # then the source's token-major route order. Every rank can rebuild all + # source ranks' inputs because they derive from deterministic seeds. + expected_rows = [] + expected_metadata = [] + for local_expert in range(experts_per_rank): + global_expert = begin + local_expert + for src_rank in range(world_size): + src_generator = torch.Generator().manual_seed(2000 + src_rank) + src_token_count = 3 + src_rank + src_activation = torch.randn(src_token_count, hidden, generator=src_generator) + for token in range(src_token_count): + for slot, slot_expert in enumerate( + ( + (src_rank + token) % experts, + (src_rank + token + 2) % experts, + ) + ): + if slot_expert == global_expert: + expected_rows.append(src_activation[token] @ global_fc1[global_expert]) + expected_metadata.append([local_expert, src_rank, token, slot]) + if expected_rows: + expected_fc1_c = torch.stack(expected_rows).to(torch.bfloat16) + else: + expected_fc1_c = torch.empty((0, 2 * intermediate), dtype=torch.bfloat16) + assert fc1_c.shape == expected_fc1_c.shape + assert route_metadata.dtype == torch.int32 + assert route_metadata.cpu().tolist() == expected_metadata + + # Backward: gradients re-dispatch along the same routes. Expected + # values come from autograd replicas over every rank's reconstructed + # inputs: token/router grads are local to this rank's tokens, while + # weight grads accumulate contributions from all source ranks. + grad_output = torch.randn(token_count, hidden, generator=torch.Generator().manual_seed(3000 + rank)) + grad_activation, grad_fc1, grad_fc2, grad_topk = op.backward( + grad_output.to(device), + activation.to(device), + global_fc1[begin:end].contiguous().to(device), + global_fc2[begin:end].contiguous().to(device), + topk_idx.to(device), + topk_weights.to(device), + fc1_c, + route_metadata, + ) + + expected_g1_global = torch.zeros_like(global_fc1) + expected_g2_global = torch.zeros_like(global_fc2) + for src_rank in range(world_size): + src_token_count = 3 + src_rank + src_activation = torch.randn(src_token_count, hidden, generator=torch.Generator().manual_seed(2000 + src_rank)) + src_topk_idx = torch.tensor( + [[(src_rank + token) % experts, (src_rank + token + 2) % experts] for token in range(src_token_count)], + dtype=torch.int64, + ) + src_topk_weights = torch.tensor([[0.625, 0.375]]).expand(src_token_count, -1).contiguous() + src_grad_output = torch.randn(src_token_count, hidden, generator=torch.Generator().manual_seed(3000 + src_rank)) + src_ga, src_g1, src_g2, src_gw = _autograd_replica_grads( + src_activation, + global_fc1, + global_fc2, + src_topk_idx, + src_topk_weights, + src_grad_output, + apply_topk_in_fc1=True, + ) + expected_g1_global += src_g1 + expected_g2_global += src_g2 + if src_rank == rank: + expected_ga, expected_gw = src_ga, src_gw + + # GPU/CPU GEMM rounding and rare bf16-stash boundary flips justify the + # loose tolerance. + torch.testing.assert_close(grad_activation.cpu(), expected_ga, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(grad_topk.cpu(), expected_gw, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(grad_fc1.cpu(), expected_g1_global[begin:end], rtol=2e-2, atol=2e-2) + torch.testing.assert_close(grad_fc2.cpu(), expected_g2_global[begin:end], rtol=2e-2, atol=2e-2) + torch.testing.assert_close(fc1_c.cpu(), expected_fc1_c) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_nccl_available() or torch.cuda.device_count() < 4, + reason="requires NCCL and at least 4 GPUs", +) +def test_four_rank_expert_parallel(tmp_path): + """One GPU per rank; NCCL runs the token all-to-all device-to-device.""" + + init_file = tmp_path / "moe_ep_nccl_init" + mp.spawn(_distributed_worker, args=(4, str(init_file)), nprocs=4, join=True) From b7138037e3c04c19c8868fbbe66dc9c4a1a2d586 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Tue, 14 Jul 2026 11:06:41 -0700 Subject: [PATCH 2/3] Update moe_ep design doc for the backward reference The doc predated the training-integration work: it still listed "no backward API" and "no auxiliary FC1 capture" as first-version boundaries and described the reference as inference-forward only. Refresh the status line, return-value contract, reference-scope notes, and validation matrix to match MoeEpReference.backward, the generate_c=True (output, fc1_c, route_metadata) return, and the current test suite. Co-Authored-By: Claude Fable 5 --- docs/fe-oss-apis/moe_ep.md | 54 +++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 21 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 7d2b57d55..211839a8b 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -1,16 +1,14 @@ # MoE + Expert Parallel API proposal Status: public API stub, implementation proposal, and executable PyTorch -reference. The API currently allocates uninitialized output storage without -launching a device kernel. Its numerical comparison is therefore marked as a -strict expected failure under `test/python/fe_api/moe_ep`; remove that marker -when backend execution is connected. +reference for both forward and backward. The API currently allocates +uninitialized output storage without launching a device kernel. Its numerical +comparison is therefore marked as a strict expected failure under +`test/python/fe_api/moe_ep`; remove that marker when backend execution is +connected. -This design is based on `mxfp8_glu_mega_instro.html`. It preserves that -kernel's routing, contiguous expert placement, SwiGLU, optional early router -weight, internal per-top-k combine plane, and final top-k reduction. It removes -workspace pointers, peer pointer mappers, streams, and scheduler tuning from the -semantic user interface. +The design removes workspace pointers, peer pointer mappers, streams, and +scheduler tuning from the semantic user interface. ## Decision summary @@ -25,9 +23,9 @@ semantic user interface. - The first half of FC1 is `gate`; the second half is `up`. SwiGLU is `silu(gate) * up`. - `output_format` means the public, post-top-k-reduction `(T, H)` result. This is - an extension of the HTML interface, whose public result is BF16. + an extension of the MegaMoE kernel interface, whose public result is BF16. - `combine_format` independently describes each per-route FC2 contribution on - the EP return path. This corresponds to the HTML's internal `combine_quant` + the EP return path. This corresponds to the `combine_quant` and `combine_sf` planes. - BF16, MXFP8, and NVFP4 are supported for both choices. Quantized output is a data-plus-scale object; it is never represented as a scale-free PyTorch @@ -134,6 +132,10 @@ capacity-based route dropping are outside this first API. `logical_shape`, and the scaled axis. `dequantize()` reconstructs a regular tensor. +With `generate_c=True`, the call instead returns +`(output, fc1_c, route_metadata)`; see the training-integration sections +below. + The return type is fixed by the constructor, so an individual module instance does not change its output structure across calls. @@ -185,7 +187,7 @@ public encoding is applied after reduction. Moving the router weight across FC2 is algebraically equivalent in exact arithmetic, but not necessarily after low-precision rounding. The option is -therefore semantic and is fixed in the constructor, matching the HTML kernel. +therefore semantic and is fixed in the constructor, matching the MegaMoE kernel. ## Block-scaled representation @@ -272,9 +274,9 @@ inbound assignments for its documented capacity policy. The conservative bound is `ep_size * max_tokens_per_rank * top_k`; a smaller bound requires an explicit router capacity/drop contract. -## Mapping to the MegaMoE HTML interface +## Mapping to the MegaMoE interface -| HTML concept | Proposed API | +| concept | Proposed API | |---|---| | `static_expert_shape=(E, 2I, H)` | `num_experts`, `intermediate_size`, `hidden_size` | | `world_size` | inferred from `ep_group` | @@ -412,8 +414,9 @@ The reference is a correctness oracle, not a performance model: accumulation order. - Scale tensors are logical, not atom-swizzled. - It uses collective push communication rather than NVSHMEM pull/remote store. -- Distributed PyTorch collectives in this reference are not an autograd - implementation. The initial scope is inference forward. +- Backward is an explicit `backward` method that re-dispatches gradients with + the same collectives; the reference is not wrapped in a + `torch.autograd.Function`, so framework integration supplies that layer. - It has no expert-capacity drop policy beyond explicit `-1` routes. ## Validation and test matrix @@ -423,10 +426,18 @@ The accompanying tests cover: - MXFP8 and NVFP4 quantize/dequantize shape and dtype contracts; - quantization along the weight reduction axis; - BF16 output for both router-weight locations and gate/up clamp; -- MXFP8 and NVFP4 public outputs; +- MXFP8 and NVFP4 quantized combine round-trips; - mixed block-scaled activation, FC1, and FC2 inputs; -- invalid expert IDs; -- two EP ranks with unequal local token counts and cross-rank routes. +- `fc1_c`/`route_metadata` capture: values, expert grouping, and row order; +- backward against autograd replicas — fp32 and block-scaled inputs, both + router-weight locations, gate/up clamp, and zero gradient for `-1` routes; +- four EP ranks (one GPU each, NCCL) with unequal local token counts and + cross-rank routes: forward, `fc1_c`/`route_metadata`, and backward gradients; +- API acceptance of block-scaled inputs and the `generate_c` allocation + contract; +- strict expected-failure gates comparing the API against the reference for + every output format, quantized inputs, and `generate_c`, armed until the + device kernel lands. Production validation should add CUDA/NVSHMEM runs for every pair of `combine_format` and `output_format`, skewed/all-to-one routing, empty experts, @@ -443,6 +454,7 @@ more surface area: - no shared/dense expert inside this operator; - no implicit top-k normalization; - no capacity factor or implicit route drop; -- no backward API; -- no auxiliary FC1 capture; +- backward semantics are fixed by `MoeEpReference.backward` (consuming the + `generate_c=True` stash); a fused device backward kernel and the + `torch.autograd` wrapper are not part of this first version; - logical scales at the Python boundary, backend swizzle internally. From 50df079e8e0f954e1d1eaca5568c7eb47e31d905 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Tue, 14 Jul 2026 11:13:04 -0700 Subject: [PATCH 3/3] Document and stub the MoeEp.backward call contract The proposed-API section only showed the constructor and __call__; add the backward signature, its argument/return tables, and the collective requirement. Give MoeEp a matching allocation-only backward stub (validating the generate_c stash shapes) with a passing allocation test and a strict-xfail gate against MoeEpReference.backward. Co-Authored-By: Claude Fable 5 --- docs/fe-oss-apis/moe_ep.md | 38 +++++++++++ python/cudnn/moe_ep/api.py | 51 ++++++++++++++ test/python/fe_api/moe_ep/test_moe_ep.py | 84 ++++++++++++++++++++++++ 3 files changed, 173 insertions(+) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 211839a8b..3c3bab54d 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -71,6 +71,18 @@ class MoeEp: | BlockScaledTensor | tuple[Tensor | BlockScaledTensor, Tensor, Tensor] ): ... + + def backward( + self, + grad_output: Tensor, + activation: Tensor | BlockScaledTensor, + fc1_weight: Tensor | BlockScaledTensor, + fc2_weight: Tensor | BlockScaledTensor, + topk_idx: Tensor, + topk_weights: Tensor, + fc1_c: Tensor, + route_metadata: Tensor, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... ``` An initialized `ep_group` enables EP. `None` deliberately means a one-rank @@ -139,6 +151,32 @@ below. The return type is fixed by the constructor, so an individual module instance does not change its output structure across calls. +### Backward call contract + +`backward` requires the operator to be constructed with `generate_c=True`; it +consumes the forward stash and is executable today as +`MoeEpReference.backward` (the device implementation is pending, like +forward). It is a collective: every rank in `ep_group` must call it, because +gradients re-dispatch along the identical forward routes. + +| Argument | Shape | Provided by | +|---|---:|---| +| `grad_output` | `(T, H)` | incoming gradient of the *dequantized* public output (all encodes are straight-through). | +| `activation`, `fc1_weight`, `fc2_weight`, `topk_idx`, `topk_weights` | as in forward | the framework re-supplies the same forward inputs; `topk_idx` deterministically regenerates the dispatch plan. | +| `fc1_c`, `route_metadata` | `(local_routes, 2I)`, `(local_routes, 4)` | the `generate_c=True` forward stash, passed back unchanged. | + +Returns four FP32 tensors: + +| Return | Shape | Meaning | +|---|---:|---| +| `grad_activation` | `(T, H)` | summed over this token's valid routes; gradients w.r.t. the dequantized activation values. | +| `grad_fc1_weight` | `(E_local, H, 2I)` | accumulated over every route this rank's experts processed, from all source ranks. | +| `grad_fc2_weight` | `(E_local, I, H)` | same accumulation domain as `grad_fc1_weight`. | +| `grad_topk_weights` | `(T, K)` | per-route router-weight gradient; exact zero at `-1` slots. | + +The save-set, recompute rules, and numerical conventions behind this +signature are specified in "Saved tensors for backward" below. + ### Example ```python diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 45d6b9cd1..e528b2a00 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -335,5 +335,56 @@ def __call__( route_metadata = torch.empty((local_routes, 4), dtype=torch.int32, device=device) return output, fc1_c, route_metadata + 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]: + """Allocate gradients for a future backend MoE+EP backward launch. + + Requires ``generate_c=True``; consumes the forward stash + (``fc1_c``, ``route_metadata``) plus the re-supplied forward inputs. + Returns float32 ``(grad_activation, grad_fc1_weight, grad_fc2_weight, + grad_topk_weights)``. Semantics are defined by + ``MoeEpReference.backward``. + """ + + 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 {tuple(grad_output.shape)}") + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + two_i = 2 * self.intermediate_size + if route_metadata.ndim != 2 or route_metadata.shape[1] != 4: + raise ValueError(f"route_metadata shape must be (local_routes, 4), got {tuple(route_metadata.shape)}") + if tuple(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 {tuple(fc1_c.shape)}") + + device = _tensor_device(activation) + # TODO: re-dispatch grad_output rows, recompute SwiGLU from fc1_c, + # accumulate weight gradients, and return route gradients to sources. + return ( + torch.empty((token_count, self.hidden_size), dtype=torch.float32, device=device), + torch.empty( + (self.experts_per_rank, self.hidden_size, two_i), + dtype=torch.float32, + device=device, + ), + torch.empty( + (self.experts_per_rank, self.intermediate_size, self.hidden_size), + dtype=torch.float32, + device=device, + ), + torch.empty((token_count, self.top_k), dtype=torch.float32, device=device), + ) + __all__ = ["BlockScaledTensor", "MoeEp", "MoeFormat", "MoeTensor"] diff --git a/test/python/fe_api/moe_ep/test_moe_ep.py b/test/python/fe_api/moe_ep/test_moe_ep.py index 083b46f5a..798e9e8cd 100644 --- a/test/python/fe_api/moe_ep/test_moe_ep.py +++ b/test/python/fe_api/moe_ep/test_moe_ep.py @@ -588,6 +588,90 @@ def test_moe_ep_api_generate_c_matches_reference(): torch.testing.assert_close(actual_metadata, expected_metadata, atol=0, rtol=0) +def test_moe_ep_api_backward_allocates_grads(): + """MoeEp.backward returns the four gradient allocations of the contract.""" + + from cudnn import MoeEp + + torch.manual_seed(53) + device = torch.device("cuda") + experts, tokens, hidden, intermediate, top_k = 2, 4, 32, 16, 2 + api = MoeEp( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + generate_c=True, + ) + activation = torch.randn(tokens, hidden, dtype=torch.bfloat16, device=device) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate, dtype=torch.bfloat16, device=device) + fc2_weight = torch.randn(experts, intermediate, hidden, dtype=torch.bfloat16, device=device) + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64, device=device) + topk_weights = torch.ones(tokens, top_k, device=device) + + _, fc1_c, route_metadata = api(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + grads = api.backward( + torch.randn(tokens, hidden, device=device), + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + expected_shapes = [ + (tokens, hidden), + (experts, hidden, 2 * intermediate), + (experts, intermediate, hidden), + (tokens, top_k), + ] + assert [tuple(grad.shape) for grad in grads] == expected_shapes + assert all(grad.dtype == torch.float32 for grad in grads) + assert all(grad.device.type == "cuda" for grad in grads) + + +@pytest.mark.xfail( + strict=True, + reason="MoeEp.backward only allocates gradients; the device implementation is not wired yet", +) +def test_moe_ep_api_backward_matches_reference(): + from cudnn import MoeEp + + torch.manual_seed(59) + device = torch.device("cuda") + experts, tokens, hidden, intermediate, top_k = 2, 4, 32, 16, 2 + activation = torch.randn(tokens, hidden, device=device) + fc1_weight = torch.randn(experts, hidden, 2 * intermediate, device=device) / 8 + fc2_weight = torch.randn(experts, intermediate, hidden, device=device) / 8 + topk_idx = torch.tensor([[0, 1], [1, 0], [0, -1], [1, 0]], dtype=torch.int64, device=device) + topk_weights = torch.tensor( + [[0.7, 0.3], [0.6, 0.4], [1.0, 0.0], [0.55, 0.45]], + dtype=torch.float32, + device=device, + ) + grad_output = torch.randn(tokens, hidden, device=device) + + kwargs = dict( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + generate_c=True, + ) + api = MoeEp(**kwargs) + reference = MoeEpReference(**kwargs) + forward_args = (activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + _, ref_fc1_c, ref_metadata = reference(*forward_args) + _, api_fc1_c, api_metadata = api(*forward_args) + + actual = api.backward(grad_output, *forward_args, api_fc1_c, api_metadata) + expected = reference.backward(grad_output, *forward_args, ref_fc1_c, ref_metadata) + for actual_grad, expected_grad in zip(actual, expected): + torch.testing.assert_close(actual_grad, expected_grad, atol=0, rtol=0) + + @pytest.mark.parametrize("input_format", ["mxfp8", "nvfp4"]) @pytest.mark.xfail( strict=True,